Regression in LR schedulers with metric tracking

Fixes #635
mittagessen · Aug 16, 2024 · 6fa8f79 · 6fa8f79
1 parent b96d0cd
commit 6fa8f79
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/kraken/lib/pretrain/model.py b/kraken/lib/pretrain/model.py
@@ -438,7 +438,10 @@ def lr_scheduler_step(self, scheduler, metric):
                 scheduler.step()
             # step every other scheduler epoch-wise
             elif self.trainer.is_last_batch:
-                scheduler.step()
+                if metric is None:
+                    scheduler.step()
+                else:
+                    scheduler.step(metric)
 
     def setup(self, stage: Optional[str] = None):
         # finalize models in case of appending/loading

diff --git a/kraken/lib/train.py b/kraken/lib/train.py
@@ -1114,7 +1114,10 @@ def lr_scheduler_step(self, scheduler, metric):
                 scheduler.step()
             # step every other scheduler epoch-wise
             elif self.trainer.is_last_batch:
-                scheduler.step()
+                if metric is None:
+                    scheduler.step()
+                else:
+                    scheduler.step(metric)
 
 
 def _configure_optimizer_and_lr_scheduler(hparams, params, len_train_set=None, loss_tracking_mode='max'):