From 8347a6c87e607a92fdc36938ec341d8f970c4322 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Wed, 30 Oct 2019 12:11:21 -0400 Subject: [PATCH] mem clear (#440) * mem clear * mem clear --- pytorch_lightning/trainer/trainer_io.py | 7 +++++++ 1 file changed, 7 insertions(+) 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)