Updated readme and setup.py after pull request.

johnmarktaylor91 · Oct 17, 2024 · 07ddc54 · 07ddc54
1 parent 25b061e
commit 07ddc54
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/tests/test_validation_and_visuals.py b/tests/test_validation_and_visuals.py
@@ -582,6 +582,7 @@ def test_buffer_rewrite_model():
     show_model_graph(
         model,
         model_input,
+        save_only=True,
         vis_opt="unrolled",
         vis_outpath=opj("visualization_outputs", "toy-networks", "buffer_rewrite_model_visible_nested_unrolled"),
         vis_buffer_layers=True,
@@ -3107,9 +3108,9 @@ class OneHotAutoEncoder(L.LightningModule):
         def __init__(self):
             super().__init__()
             self.model = nn.Sequential(
-                    nn.Linear(16, 4),
-                    nn.ReLU(),
-                    nn.Linear(4, 16),
+                nn.Linear(16, 4),
+                nn.ReLU(),
+                nn.Linear(4, 16),
             )
 
         def forward(self, x):
@@ -3127,4 +3128,3 @@ def forward(self, x):
         vis_outpath=opj("visualization_outputs", "lightning", "one-hot-autoencoder"),
     )
     assert validate_saved_activations(model, model_inputs)
-