2024-10-20 00:10:27 +02:00
1 changed files with 2 additions and 2 deletions
@@ -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)