fix: update example autoencoder.py to reflect args (#6638)

* fix: update example autoencoder.py to reflect args

* Update pl_examples/basic_examples/autoencoder.py

Co-authored-by: Carlos Mocholí <[email protected]>

Author: bmahlbrand

pl_examples/basic_examples/autoencoder.py

@@ -39,17 +39,17 @@ class LitAutoEncoder(pl.LightningModule):
     )
     """
 
-    def __init__(self):
+    def __init__(self, hidden_dim: int = 64):
         super().__init__()
         self.encoder = nn.Sequential(
-            nn.Linear(28 * 28, 64),
+            nn.Linear(28 * 28, hidden_dim),
             nn.ReLU(),
-            nn.Linear(64, 3),
+            nn.Linear(hidden_dim, 3),
         )
         self.decoder = nn.Sequential(
-            nn.Linear(3, 64),
+            nn.Linear(3, hidden_dim),
             nn.ReLU(),
-            nn.Linear(64, 28 * 28),
+            nn.Linear(hidden_dim, 28 * 28),
         )
 
     def forward(self, x):
@@ -94,7 +94,7 @@ def cli_main():
     # ------------
     parser = ArgumentParser()
     parser.add_argument('--batch_size', default=32, type=int)
-    parser.add_argument('--hidden_dim', type=int, default=128)
+    parser.add_argument('--hidden_dim', type=int, default=64)
     parser = pl.Trainer.add_argparse_args(parser)
     args = parser.parse_args()
 
@@ -112,7 +112,7 @@ def cli_main():
     # ------------
     # model
     # ------------
-    model = LitAutoEncoder()
+    model = LitAutoEncoder(args.hidden_dim)
 
     # ------------
     # training