Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fixes CPU and hanging GPU crash (#2118) 

* training batch clean up

* training batch clean up

* training batch clean up
This commit is contained in:
William Falcon 2020-06-08 16:30:20 -04:00 committed by GitHub
parent 9e8716afe8
commit 0bd7780adc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
pytorch_lightning/trainer/trainer.py
View File

@ -1033,7 +1033,7 @@ class Trainer(
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}'):