Handle corrupt ckpt

Noble-Lab · Nov 29, 2023 · 0e3bcb8 · 0e3bcb8
1 parent bfb168e
commit 0e3bcb8
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/casanovo/denovo/model_runner.py b/casanovo/denovo/model_runner.py
@@ -289,9 +289,18 @@ def initialize_model(self, train: bool) -> None:
                         f"using the checkpoint."
                     )
         except RuntimeError:
-            raise RuntimeError(
-                "Weights file incompatible with the current version of Casanovo."
-            )
+            # This only doesn't work if the weights are from an older version
+            try:
+                self.model = Spec2Pep.load_from_checkpoint(
+                    self.model_filename,
+                    map_location=device,
+                    **model_params,
+                )
+            except RuntimeError:
+                raise RuntimeError(
+                    "Weights file incompatible "
+                    "with the current version of Casanovo. "
+                )
 
     def initialize_data_module(
         self,