Fix progress bar updates for Pod Training

Lightning-AI · Jul 2, 2021 · 98203e1 · 98203e1
1 parent d51b0ae
commit 98203e1
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/pytorch_lightning/plugins/training_type/tpu_spawn.py b/pytorch_lightning/plugins/training_type/tpu_spawn.py
@@ -152,6 +152,9 @@ def new_process(self, process_idx: int, trainer, mp_queue) -> None:
         # set warning rank
         rank_zero_only.rank = self.global_rank
 
+        if self.tpu_global_core_rank == 0 and int(os.getenv(xenv.TPUVM_MODE, 0)) == 1:
+            print(' ', end='', flush=True)
+
         if self.tpu_global_core_rank != 0 and trainer.progress_bar_callback is not None:
             trainer.progress_bar_callback.disable()