import MPI from mpi4py

2024-08-19 01:01:54 -07:00 · 2024-08-19 01:01:54 -07:00 · 5feca7f4b6
parent 64f86b7aa6
commit 5feca7f4b6
1 changed files with 2 additions and 2 deletions
--- a/seem/utils/distributed.py
+++ b/seem/utils/distributed.py
@ -3,7 +3,7 @@ import time
 import torch
 import pickle
 import subprocess
-
+from mpi4py import MPI
 import torch.distributed as dist


@ -108,4 +108,4 @@ def synchronize():

    _send_and_wait(0)
    # now sync on the main process
-    _send_and_wait(1)
+    _send_and_wait(1)