diff --git a/CollaborativeCoding/dataloaders/download.py b/CollaborativeCoding/dataloaders/download.py
index 08e7b27..e165d5f 100644
--- a/CollaborativeCoding/dataloaders/download.py
+++ b/CollaborativeCoding/dataloaders/download.py
@@ -98,11 +98,11 @@ def download_svhn(path, train: bool = True):
         train_data = parent_path / "train_32x32.mat"
         test_data = parent_path / "test_32x32.mat"
 
-        if not train_data.exists():
+        if not train_data.is_file():
             download_svhn(parent_path, train=True)
-        if not test_data.exists():
+        if not test_data.is_file():
             download_svhn(parent_path, train=False)
-
+        print(test_data)
         train_labels = loadmat(train_data)["y"]
         test_labels = loadmat(test_data)["y"]
 
diff --git a/CollaborativeCoding/models/magnus_model.py b/CollaborativeCoding/models/magnus_model.py
index 7a79057..244a5ff 100644
--- a/CollaborativeCoding/models/magnus_model.py
+++ b/CollaborativeCoding/models/magnus_model.py
@@ -2,7 +2,7 @@
 
 
 class MagnusModel(nn.Module):
-    def __init__(self, image_shape, num_classes: int, nr_channels: int = 1):
+    def __init__(self, image_shape, num_classes: int):
         """
         Initializes the MagnusModel, a neural network designed for image classification tasks.
         The model consists of three linear layers, each with 133 neurons, and uses ReLU activation
@@ -14,11 +14,11 @@ def __init__(self, image_shape, num_classes: int, nr_channels: int = 1):
             nr_channels (int): The number of channels in the input image.
         """
         super().__init__()
-        *_, H, W = image_shape
+        C, H, W = image_shape[-3], image_shape[-2], image_shape[-1]
         self.layer1 = nn.Sequential(
             *(
                 [
-                    nn.Linear(nr_channels * H * W, 133),
+                    nn.Linear(C * H * W, 133),
                     nn.ReLU(),
                 ]
             )