get from optimization config

mllam · Jan 27, 2025 · 14ab353 · 14ab353
1 parent d55bf91
commit 14ab353
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py
@@ -195,14 +195,14 @@ def configure_optimizers(self):
             self.parameters(), lr=self.optimization_config.lr, betas=(0.9, 0.95)
         )
 
-        if self.optimization_config.scheduler:
-            scheduler_class = getattr(
-                torch.optim.lr_scheduler, self.optimization_config.scheduler
+        if self.optimization_config.lr_scheduler:
+            lr_scheduler_class = getattr(
+                torch.optim.lr_scheduler, self.optimization_config.lr_scheduler
             )
-            scheduler = scheduler_class(
-                optimizer, **self.optimization_config.scheduler_kwargs
+            lr_scheduler = lr_scheduler_class(
+                optimizer, **self.optimization_config.lr_scheduler_kwargs
             )
-            return [optimizer], [scheduler]
+            return [optimizer], [lr_scheduler]
 
         return optimizer