From 6541a47668cf458c61bc7343ede4735408665752 Mon Sep 17 00:00:00 2001
From: Brian Bjarke Jensen <bbj@skov.dk>
Date: Sun, 28 Jul 2024 21:19:47 +0200
Subject: [PATCH] added option to download resnet18 weights

---
 model/src/models/resnet18_head.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/model/src/models/resnet18_head.py b/model/src/models/resnet18_head.py
index 00b06c8..e918d65 100644
--- a/model/src/models/resnet18_head.py
+++ b/model/src/models/resnet18_head.py
@@ -5,11 +5,15 @@ import torchvision
 
 
 class ResNet18Head(nn.Module):
-    def __init__(self):
+    def __init__(self, download_resnet_weights: bool = False):
         super().__init__()
         # copy out parts from ResNet18 with weights
+        if download_resnet_weights:
+            weights = torchvision.models.ResNet18_Weights.IMAGENET1K_V1
+        else:
+            weights = None
         resnet18 = torchvision.models.resnet18(
-            # weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1,
+            weights=weights,
         )
         # save relevant layers
         self.conv1 = resnet18.conv1