Bugfix: load checkpoint also if "model_state_dict" is not a key

funkelab · Sep 3, 2023 · 0c13f37 · 0c13f37
1 parent 8f8b6aa
commit 0c13f37
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/gunpowder/torch/nodes/predict.py b/gunpowder/torch/nodes/predict.py
@@ -104,7 +104,7 @@ def start(self):
             if "model_state_dict" in checkpoint:
                 self.model.load_state_dict(checkpoint["model_state_dict"])
             else:
-                self.model.load_state_dict()
+                self.model.load_state_dict(checkpoint)
 
     def predict(self, batch, request):
         inputs = self.get_inputs(batch)