fix

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>
ray-project · Nov 2, 2023 · 69756f4 · 69756f4
1 parent ddc8efe
commit 69756f4
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/python/ray/train/lightning/_lightning_utils.py b/python/ray/train/lightning/_lightning_utils.py
@@ -288,6 +288,8 @@ def on_train_epoch_end(self, trainer, pl_module) -> None:
             checkpoint = Checkpoint.from_directory(tmpdir)
         train.report(metrics=metrics, checkpoint=checkpoint)
 
+        torch.distributed.barrier()
+
         if self.local_rank == 0:
             shutil.rmtree(tmpdir)