Update SegmentationSegformer checkpoint loading with optional weights…

… parameter
NRCan · Feb 6, 2025 · 49e8929 · 49e8929
1 parent 916c76e
commit 49e8929
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/geo_deep_learning/tasks_with_models/segmentation_segformer.py b/geo_deep_learning/tasks_with_models/segmentation_segformer.py
@@ -160,7 +160,9 @@ def export_model(self, checkpoint_path: str, export_path: str, datamodule: Light
         map_location = "cuda"
         if self.device.type == "cpu":
             map_location = "cpu"
-        best_model = self.__class__.load_from_checkpoint(checkpoint_path, map_location=map_location)
+        best_model = self.__class__.load_from_checkpoint(checkpoint_path,
+                                                         weights_from_checkpoint_path=None,
+                                                         map_location=map_location)
         best_model.eval()
 
         scrpted_model = script_model(best_model.model, datamodule, self.num_classes, from_logits=True)