wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-09 11:32:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1645 from danpcampbell/bugfix/1637_slurm-fix

Browse Source 
Only call load_spawn_weights if COLAB_GPU or KAGGLE_URL_BASE
This commit is contained in:
William Falcon
2020-04-27 20:08:06 -04:00
committed by GitHub
parent b0ec7a0655 981758fd04
commit e40e27f59e
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/trainer/trainer.py
+3 -2
View File
@@ -750,8 +750,9 @@ class Trainer(
# train
mp.spawn(self.ddp_train, nprocs=self.num_processes, args=(model,))
# load weights if not interrupted
self.load_spawn_weights(model)
self.model = model
if os.getenv('COLAB_GPU') or os.getenv('KAGGLE_URL_BASE'):
self.load_spawn_weights(model)
self.model = model
# 1 gpu or dp option triggers training using DP module
# easier to avoid NCCL issues