Minor fix to loading models with incompatible layers (#641)

computational-cell-analytics · Jun 26, 2024 · 5e16964 · 5e16964
1 parent 8e9750c
commit 5e16964
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/micro_sam/util.py b/micro_sam/util.py
@@ -409,10 +409,11 @@ def _handle_checkpoint_loading(sam, model_state):
     reference_state = sam.state_dict()
 
     for k, v in model_state.items():
-        if reference_state[k].size() == v.size():
-            new_state_dict[k] = v
-        else:
-            mismatched_layers.append(k)
+        if k in reference_state:  # This is done to get rid of unwanted layers from pretrained SAM.
+            if reference_state[k].size() == v.size():
+                new_state_dict[k] = v
+            else:
+                mismatched_layers.append(k)
 
     reference_state.update(new_state_dict)