Fix model loading issue with new lightning versions

teamtomo · Oct 30, 2023 · 26db9de · 26db9de
1 parent c132a79
commit 26db9de
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/membrain_seg/segmentation/segment.py b/src/membrain_seg/segmentation/segment.py
@@ -78,8 +78,7 @@ def segment(
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     # Initialize the model and load trained weights from checkpoint
-    pl_model = SemanticSegmentationUnet()
-    pl_model = pl_model.load_from_checkpoint(model_checkpoint, map_location=device)
+    pl_model = SemanticSegmentationUnet.load_from_checkpoint(model_checkpoint, map_location=device, strict=False)
     pl_model.to(device)
 
     # Preprocess the new data