added option to download weights

2024-07-28 21:25:24 +02:00
parent 6541a47668
commit be6833ac6a
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()