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

Only call load_spawn_weights if COLAB_GPU or KAGGLE_URL_BASE

Browse Source 
environment variables are set
This commit is contained in:
Dan Campbell
2020-04-27 19:46:29 -04:00
parent b0ec7a0655
commit 981758fd04
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