From 0bd7780adc4d68007946cf380a6a24e1a08d99d1 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Mon, 8 Jun 2020 16:30:20 -0400 Subject: [PATCH] Fixes CPU and hanging GPU crash (#2118) * training batch clean up * training batch clean up * training batch clean up --- pytorch_lightning/trainer/trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index aff19fb253f38..1f5b73f9be364 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -1033,7 +1033,7 @@ def run_pretrain_routine(self, model: LightningModule): self.early_stop_callback._validate_condition_metric(callback_metrics) # clear cache before training - if self.on_gpu: + if self.on_gpu and self.root_gpu is not None: # use context because of: # https://discuss.pytorch.org/t/out-of-memory-when-i-use-torch-cuda-empty-cache/57898 with torch.cuda.device(f'cuda:{self.root_gpu}'):