added option to download weights
This commit is contained in:
@@ -35,10 +35,10 @@ from .visual_syntax import VisualSyntaxTail
|
|||||||
class VisualCommunicationModel(nn.Module):
|
class VisualCommunicationModel(nn.Module):
|
||||||
"""Visual communication model."""
|
"""Visual communication model."""
|
||||||
|
|
||||||
def __init__(self):
|
def __init__(self, download_resnet_weights: bool = False):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
# store other models
|
# store other models
|
||||||
self.resnet_head = ResNet18Head()
|
self.resnet_head = ResNet18Head(download_resnet_weights)
|
||||||
self.visual_syntax_tail = VisualSyntaxTail()
|
self.visual_syntax_tail = VisualSyntaxTail()
|
||||||
self.contact_tail = ContactTail()
|
self.contact_tail = ContactTail()
|
||||||
self.angle_tail = AngleTail()
|
self.angle_tail = AngleTail()
|
||||||
|
|||||||
Reference in New Issue
Block a user