Update nnUNetTrainer.py

Author: SimoneBendazzoli93

nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py

@@ -1209,7 +1209,8 @@ def load_checkpoint(self, filename_or_checkpoint: Union[dict, str]) -> None:
                 self.network._orig_mod.load_state_dict(new_state_dict)
             else:
                 self.network.load_state_dict(new_state_dict)
-        self.optimizer.load_state_dict(checkpoint['optimizer_state'])
+        if 'optimizer_state' in checkpoint.keys():
+            self.optimizer.load_state_dict(checkpoint['optimizer_state'])
         if self.grad_scaler is not None:
             if checkpoint['grad_scaler_state'] is not None:
                 self.grad_scaler.load_state_dict(checkpoint['grad_scaler_state'])