Made a small mistake in models code

guydav · Nov 28, 2022 · 47b70ae · 47b70ae
1 parent 807a2c9
commit 47b70ae
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/simple_relational_reasoning/embeddings/models.py b/simple_relational_reasoning/embeddings/models.py
@@ -183,9 +183,9 @@ def build_model(model_name: str, device: str, pretrained: bool = True, saycam: t
             model = models.resnext50_32x4d(weights=weights)
             model.fc_backup = model.fc
             model.embedding_dim = MODEL_EMBEDDING_DIMENSIONS[(model_name, unpooled_output)]  # type: ignore
-            model.module.fc = nn.Identity()  # type: ignore
+            model.fc = nn.Identity()  # type: ignore
             if unpooled_output:
-                model.module.avgpool = nn.Identity()  # type: ignore
+                mode.avgpool = nn.Identity()  # type: ignore
             model = model.to(device)
 
     if model is None: