diff --git a/model/src/models/visual_communication.py b/model/src/models/visual_communication.py
index e2b9d8c..a0b55e1 100644
--- a/model/src/models/visual_communication.py
+++ b/model/src/models/visual_communication.py
@@ -35,10 +35,10 @@ from .visual_syntax import VisualSyntaxTail
 class VisualCommunicationModel(nn.Module):
     """Visual communication model."""
 
-    def __init__(self):
+    def __init__(self, download_resnet_weights: bool = False):
         super().__init__()
         # store other models
-        self.resnet_head = ResNet18Head()
+        self.resnet_head = ResNet18Head(download_resnet_weights)
         self.visual_syntax_tail = VisualSyntaxTail()
         self.contact_tail = ContactTail()
         self.angle_tail = AngleTail()