feat(distributed): add barrier

metaopt · XuehaiPan · Sep 28, 2022 · Sep 25, 2022 · Sep 26, 2022 · Sep 26, 2022
commit dfcbe416aceaea39dfbef63b75f69caec4a6d5aa
diff --git a/torchopt/distributed/world.py b/torchopt/distributed/world.py
@@ -164,6 +164,7 @@ def wrapped(*args, **kwargs):
             )
             atexit.register(rpc.shutdown, graceful=True)
             if worker_init_fn is not None:
+                barrier()
                 worker_init_fn()
             barrier()
             return func(*args, **kwargs)