From 981758fd04934b58545b4426c0beb548ebdb8c17 Mon Sep 17 00:00:00 2001
From: Dan Campbell <dpc@delphiki-consulting.com>
Date: Mon, 27 Apr 2020 19:46:29 -0400
Subject: [PATCH] Only call load_spawn_weights if COLAB_GPU or KAGGLE_URL_BASE
 environment variables are set

---
 pytorch_lightning/trainer/trainer.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
index deed0098..7bfd97bf 100644
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -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