test load_spawn_weights

Lightning-AI · Aug 14, 2020 · 8f83bd2 · 8f83bd2
1 parent 056bf5e
commit 8f83bd2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_lightning/accelerators/tpu_backend.py b/pytorch_lightning/accelerators/tpu_backend.py
@@ -93,7 +93,7 @@ def __load_weights_on_main_process(self):
 
         # load weights if not interrupted
         if self.trainer.on_colab_kaggle and not self.trainer.testing:
-            self.trainer.load_spawn_weights(model)
+            model = self.trainer.load_spawn_weights(model)
 
         self.trainer.model = model