Remove default input shape (#21)

* remove default input shape * add input shape to the tests
willyfh · Feb 24, 2024 · 534bb3e · 534bb3e
1 parent aee21c0
commit 534bb3e
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/tests/test_layered.py b/tests/test_layered.py
@@ -55,20 +55,20 @@ def forward(self, x):
 
 def test_sequential_model_layered_view_runs(sequential_model):
     try:
-        _ = layered_view(sequential_model)
+        _ = layered_view(sequential_model, input_shape=(1, 3, 224, 224))
     except Exception as e:
         pytest.fail(f"layered_view raised an exception with Sequential model: {e}")
 
 
 def test_module_list_model_layered_view_runs(module_list_model):
     try:
-        _ = layered_view(module_list_model)
+        _ = layered_view(module_list_model, input_shape=(1, 3, 224, 224))
     except Exception as e:
         pytest.fail(f"layered_view raised an exception with ModuleList model: {e}")
 
 
 def test_custom_model_layered_view_runs(custom_model):
     try:
-        _ = layered_view(custom_model)
+        _ = layered_view(custom_model, input_shape=(1, 3, 224, 224))
     except Exception as e:
         pytest.fail(f"layered_view raised an exception with Custom model: {e}")
diff --git a/visualtorch/layered.py b/visualtorch/layered.py
@@ -24,7 +24,7 @@
 
 def layered_view(
     model,
-    input_shape=(1, 3, 224, 224),
+    input_shape,
     to_file: str | None = None,
     min_z: int = 10,
     min_xy: int = 10,