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