diff --git a/multitalent/training/lr_scheduler/polylr.py b/multitalent/training/lr_scheduler/polylr.py
index 2915fa5..af093e1 100644
--- a/multitalent/training/lr_scheduler/polylr.py
+++ b/multitalent/training/lr_scheduler/polylr.py
@@ -8,7 +8,7 @@ def __init__(self, optimizer, initial_lr: float, max_steps: int, exponent: float
         self.max_steps = max_steps
         self.exponent = exponent
         self.ctr = 0
-        super().__init__(optimizer, current_step if current_step is not None else -1, False)
+        super().__init__(optimizer, current_step if current_step is not None else -1)
 
     def step(self, current_step=None):
         if current_step is None or current_step == -1: