From be6833ac6a7640dfe53fcdaa7c16bebc96f11d1f Mon Sep 17 00:00:00 2001 From: Brian Bjarke Jensen Date: Sun, 28 Jul 2024 21:25:24 +0200 Subject: [PATCH] added option to download weights --- model/src/models/visual_communication.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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()