Fixes CPU and hanging GPU crash (#2118)
* training batch clean up * training batch clean up * training batch clean up
This commit is contained in:
parent
9e8716afe8
commit
0bd7780adc
|
@ -1033,7 +1033,7 @@ class Trainer(
|
||||||
self.early_stop_callback._validate_condition_metric(callback_metrics)
|
self.early_stop_callback._validate_condition_metric(callback_metrics)
|
||||||
|
|
||||||
# clear cache before training
|
# clear cache before training
|
||||||
if self.on_gpu:
|
if self.on_gpu and self.root_gpu is not None:
|
||||||
# use context because of:
|
# use context because of:
|
||||||
# https://discuss.pytorch.org/t/out-of-memory-when-i-use-torch-cuda-empty-cache/57898
|
# 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}'):
|
with torch.cuda.device(f'cuda:{self.root_gpu}'):
|
||||||
|
|
Loading…
Reference in New Issue