Examples: expose learning rate (Lightning-AI#17513)

Co-authored-by: thomas <thomas@thomass-MacBook-Pro.local>
SkafteNicki · Apr 28, 2023 · cafdc6d · cafdc6d
1 parent 3867045
commit cafdc6d
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/examples/pytorch/basics/autoencoder.py b/examples/pytorch/basics/autoencoder.py
@@ -114,8 +114,9 @@ class LitAutoEncoder(LightningModule):
     )
     """
 
-    def __init__(self, hidden_dim: int = 64):
+    def __init__(self, hidden_dim: int = 64, learning_rate=10e-3):
         super().__init__()
+        self.save_hyperparameters()
         self.encoder = nn.Sequential(nn.Linear(28 * 28, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 3))
         self.decoder = nn.Sequential(nn.Linear(3, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 28 * 28))
 
@@ -138,7 +139,7 @@ def predict_step(self, batch, batch_idx, dataloader_idx=None):
         return self(x)
 
     def configure_optimizers(self):
-        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
         return optimizer
 
     def _prepare_batch(self, batch):