diff --git a/model/src/models/resnet18_head.py b/model/src/models/resnet18_head.py
index e647fb5..00b06c8 100644
--- a/model/src/models/resnet18_head.py
+++ b/model/src/models/resnet18_head.py
@@ -9,7 +9,7 @@ class ResNet18Head(nn.Module):
         super().__init__()
         # copy out parts from ResNet18 with weights
         resnet18 = torchvision.models.resnet18(
-            weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1,
+            # weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1,
         )
         # save relevant layers
         self.conv1 = resnet18.conv1
@@ -21,7 +21,7 @@ class ResNet18Head(nn.Module):
         self.layer3 = resnet18.layer3
         self.layer4 = resnet18.layer4
         self.avgpool = resnet18.avgpool
-        self.flat = nn.Flatten()    # size 512
+        self.flat = nn.Flatten()  # size 512
 
     def forward(self, x):
         x = self.conv1(x)