diff --git a/pytorch_lightning/trainer/trainer_io.py b/pytorch_lightning/trainer/trainer_io.py index b8de2bbf..dd38c051 100644 --- a/pytorch_lightning/trainer/trainer_io.py +++ b/pytorch_lightning/trainer/trainer_io.py @@ -32,10 +32,17 @@ class TrainerIOMixin(object): :param model: :return: """ + # clear cache before restore + if self.on_gpu: + torch.cuda.empty_cache() # if script called from hpc resubmit, load weights did_restore_hpc_weights = self.restore_hpc_weights_if_needed(model) + # clear cache after restore + if self.on_gpu: + torch.cuda.empty_cache() + if not did_restore_hpc_weights: # restore weights if same exp version self.restore_state_if_checkpoint_exists(model)