Updated bs-scheduler

ancestor-mithril · Dec 6, 2024 · 69abb73 · 69abb73
1 parent 3213e57
commit 69abb73
Showing 1 changed file with 1 addition and 7 deletions.
diff --git a/utils/trainer.py b/utils/trainer.py
@@ -243,13 +243,7 @@ def scheduler_step(self, metrics):
             "ReduceLROnPlateau",
             "IncreaseBSOnPlateau",
         ):
-            scheduler_metric = metrics[self.scheduler_metric]
-
-            if type(self.scheduler).__name__ == "IncreaseBSOnPlateau":
-                self.scheduler.step(metric=scheduler_metric)
-                # TODO: keep the same name
-            else:
-                self.scheduler.step(scheduler_metric)
+            self.scheduler.step(metrics=metrics[self.scheduler_metric])
         else:
             self.scheduler.step()