From fc26078e395f8a001f4c6dd7b3fe7ca202f914a3 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Tue, 30 Jun 2020 17:20:18 -0400 Subject: [PATCH] fix tpu tests --- pytorch_lightning/trainer/trainer.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index 7c8a89bffb87a..c97dbcfa2808f 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -946,7 +946,9 @@ def fit( xmp.spawn(self.tpu_train, args=(model,), nprocs=self.tpu_cores, start_method=start_method) # load weights if not interrupted - self.load_spawn_weights(model) + if self.on_colab_kaggle: + self.load_spawn_weights(model) + self.model = model # ON CPU