fixes for compatibility with new torch versions #5

Author: ykirchhoff

nnunetv2/inference/predict_from_raw_data.py

@@ -82,7 +82,7 @@ def initialize_from_trained_model_folder(self, model_training_output_dir: str,
         for i, f in enumerate(use_folds):
             f = int(f) if f != 'all' else f
             checkpoint = torch.load(join(model_training_output_dir, f'fold_{f}', checkpoint_name),
-                                    map_location=torch.device('cpu'))
+                                    map_location=torch.device('cpu'), weights_only=False)
             if i == 0:
                 trainer_name = checkpoint['trainer_name']
                 configuration_name = checkpoint['init_args']['configuration']

nnunetv2/run/load_pretrained_weights.py

@@ -17,9 +17,9 @@ def load_pretrained_weights(network, fname, verbose=False):
 
     """
     if dist.is_initialized():
-        saved_model = torch.load(fname, map_location=torch.device('cuda', dist.get_rank()))
+        saved_model = torch.load(fname, map_location=torch.device('cuda', dist.get_rank()), weights_only=False)
     else:
-        saved_model = torch.load(fname)
+        saved_model = torch.load(fname, weight_only=False)
     pretrained_dict = saved_model['network_weights']
 
     skip_strings_in_pretrained = [

nnunetv2/training/lr_scheduler/polylr.py

@@ -8,7 +8,7 @@ def __init__(self, optimizer, initial_lr: float, max_steps: int, exponent: float
         self.max_steps = max_steps
         self.exponent = exponent
         self.ctr = 0
-        super().__init__(optimizer, current_step if current_step is not None else -1, False)
+        super().__init__(optimizer, current_step if current_step is not None else -1)
 
     def step(self, current_step=None):
         if current_step is None or current_step == -1:

nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py

@@ -1181,7 +1181,7 @@ def load_checkpoint(self, filename_or_checkpoint: Union[dict, str]) -> None:
             self.initialize()
 
         if isinstance(filename_or_checkpoint, str):
-            checkpoint = torch.load(filename_or_checkpoint, map_location=self.device)
+            checkpoint = torch.load(filename_or_checkpoint, map_location=self.device, weights_only=False)
         # if state dict comes from nn.DataParallel but we use non-parallel model here then the state dict keys do not
         # match. Use heuristic to make it match
         new_state_dict = {}