Make use of learning rate scheduler optional (ashleve#449)

* Made `trainer.configure_optimizers()` robust to unspecified learning rate schedulers
razvanursu · Oct 5, 2022 · ad0f46c · ad0f46c
1 parent 89f9f02
commit ad0f46c
Showing 1 changed file with 12 additions and 11 deletions.
diff --git a/src/models/mnist_module.py b/src/models/mnist_module.py
@@ -124,17 +124,18 @@ def configure_optimizers(self):
             https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#configure-optimizers
         """
         optimizer = self.hparams.optimizer(params=self.parameters())
-        scheduler = self.hparams.scheduler(optimizer=optimizer)
-
-        return {
-            "optimizer": optimizer,
-            "lr_scheduler": {
-                "scheduler": scheduler,
-                "monitor": "val/loss",
-                "interval": "epoch",
-                "frequency": 1,
-            },
-        }
+        if self.hparams.scheduler is not None:
+            scheduler = self.hparams.scheduler(optimizer=optimizer)
+            return {
+                "optimizer": optimizer,
+                "lr_scheduler": {
+                    "scheduler": scheduler,
+                    "monitor": "val/loss",
+                    "interval": "epoch",
+                    "frequency": 1,
+                },
+            }
+        return {"optimizer": optimizer}
 
 
 if __name__ == "__main__":