fix: fix bug in loading statedict

instadeepai · Jan 10, 2025 · 536e118 · 536e118
1 parent 02f81f2
commit 536e118
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/instageo/model/model.py b/instageo/model/model.py
@@ -160,7 +160,12 @@ def __init__(
         if freeze_backbone:
             for param in model.parameters():
                 param.requires_grad = False
-        _ = model.load_state_dict(checkpoint, strict=False)
+        filtered_checkpoint_state_dict = {
+            key[len("encoder.") :]: value
+            for key, value in checkpoint.items()
+            if key.startswith("encoder.")
+        }
+        _ = model.load_state_dict(filtered_checkpoint_state_dict)
 
         self.prithvi_100M_backbone = model