Merge branch 'main' into aimseg-data

constantinpape · Feb 22, 2025 · ccc452a · ccc452a
2 parents 3766d37 + 9cc87d2
commit ccc452a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_em/trainer/default_trainer.py b/torch_em/trainer/default_trainer.py
@@ -278,7 +278,7 @@ def load_logger(self, kwarg_name: str, optional: bool):
     def _get_save_dict(save_path, device):
         if not os.path.exists(save_path):
             raise ValueError(f"Cannot find checkpoint {save_path}")
-        return torch.load(save_path, map_location=device)
+        return torch.load(save_path, map_location=device, weights_only=False)
 
     @classmethod
     def from_checkpoint(