MagnusModel fix

Author: Seilmast

.gitignore

@@ -11,7 +11,7 @@ wandb_api.py
 #Magnus specific
 docker/*
 job*
-
+env2/*
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
@@ -150,6 +150,7 @@ ENV/
 env.bak/
 venv.bak/
 
+
 # Spyder project settings
 .spyderproject
 .spyproject

utils/models/magnus_model.py

@@ -17,7 +17,7 @@ def __init__(self, image_shape, num_classes: int, nr_channels: int):
             MagnusModel (nn.Module): An instance of the MagnusModel neural network.
         """
         super().__init__()
-        _, H, W = image_shape
+        *_, H, W = image_shape
 
         self.layer1 = nn.Sequential(
             *(
@@ -55,12 +55,12 @@ def forward(self, x):
 if __name__ == "__main__":
     import torch as th
 
-    data_shape = [28, 28]
+    image_shape = (3, 28, 28)
+    n, c, h, w = 5, *image_shape
+    model = MagnusModel([h, w], 10, c)
 
-    data_shape = (3, *data_shape)
-    model = MagnusModel(data_shape, 10)
-
-    dummy_img = th.rand((5, *data_shape))
-    print(dummy_img.size())
+    x = th.rand((n, c, h, w))
     with th.no_grad():
-        print(model(dummy_img).size())
+        y = model(x)
+
+    assert y.shape == (n, 10), f"Shape: {y.shape}"