From 0f1b7208e664366373cc9a4418eb822c40ef4e80 Mon Sep 17 00:00:00 2001
From: Sebastian Hoffmann <shoffmann.git@gmail.com>
Date: Tue, 2 Apr 2024 18:48:12 +0200
Subject: [PATCH] feat: synchronize more

---
 dmlcloud/stage.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/dmlcloud/stage.py b/dmlcloud/stage.py
index 6ace04d..d4b5a25 100644
--- a/dmlcloud/stage.py
+++ b/dmlcloud/stage.py
@@ -136,6 +136,7 @@ def run(self):
             if self._stop_requested:
                 break
         self._post_stage()
+        dist.barrier()  # this will time out if not all processes reach this point
 
     def _pre_stage(self):
         self.start_time = datetime.now()