add tests for single scalar return from training

Lightning-AI · Jul 11, 2020 · 5dce449 · 5dce449
1 parent a8f28ba
commit 5dce449
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/pytorch_lightning/trainer/distrib_parts.py b/pytorch_lightning/trainer/distrib_parts.py
@@ -192,6 +192,8 @@ def tpu_train(self, tpu_core_idx, model):
             model.prepare_data()
             self._is_data_prepared = True
 
+        self.barrier()
+
         # call setup after the ddp process has connected
         self.setup('fit')
         if self.is_function_implemented('setup', model):