Merge regression with LR schedulers

mittagessen · Dec 3, 2023 · baff962 · baff962
1 parent ba4c81f
commit baff962
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/kraken/lib/train.py b/kraken/lib/train.py
@@ -665,7 +665,7 @@ def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
             for pg in optimizer.param_groups:
                 pg["lr"] = lr_scale * self.hparams.hyper_params['lrate']
 
-    def lr_scheduler_step(self, scheduler, optimizer_idx, metric):
+    def lr_scheduler_step(self, scheduler, metric):
         if not self.hparams.hyper_params['warmup'] or self.trainer.global_step >= self.hparams.hyper_params['warmup']:
             # step OneCycleLR each batch if not in warmup phase
             if isinstance(scheduler, lr_scheduler.OneCycleLR):
@@ -1080,7 +1080,7 @@ def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
             for pg in optimizer.param_groups:
                 pg["lr"] = lr_scale * self.hparams.hyper_params['lrate']
 
-    def lr_scheduler_step(self, scheduler, optimizer_idx, metric):
+    def lr_scheduler_step(self, scheduler, metric):
         if not self.hparams.hyper_params['warmup'] or self.trainer.global_step >= self.hparams.hyper_params['warmup']:
             # step OneCycleLR each batch if not in warmup phase
             if isinstance(scheduler, lr_scheduler.OneCycleLR):