Fix load_checkpoint encoder prefix (#290)

computational-cell-analytics · Dec 20, 2023 · d736b77 · d736b77
1 parent cc4fe4b
commit d736b77
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/micro_sam/training/joint_sam_trainer.py b/micro_sam/training/joint_sam_trainer.py
@@ -41,9 +41,10 @@ def load_checkpoint(self, checkpoint="best"):
         # let's get the image encoder params from sam
         sam_state = save_dict["model_state"]
         encoder_state = []
+        prune_prefix = "sam.image_"
         for k, v in sam_state.items():
-            if k.startswith("image_encoder"):
-                encoder_state.append((k, v))
+            if k.startswith(prune_prefix):
+                encoder_state.append((k[len(prune_prefix):], v))
         encoder_state = OrderedDict(encoder_state)
 
         # let's get the decoder params from unetr