Fix MPI library issue

app.py

@@ -9,8 +9,8 @@ import torch.backends.cudnn
 import admin.settings as ws_settings
 import os
 os.environ["CUDA_VISIBLE_DEVICES"] = "7"
-os.environ["OPENAI_LOGDIR"] = "./logs"
-os.environ["MPI_DISABLED"] = "1"
+# os.environ["OPENAI_LOGDIR"] = "./logs"
+# os.environ["MPI_DISABLED"] = "1"
 import torch
 import torch.distributed as dist
 import torchvision.transforms as transforms

packages.txt

@@ -0,0 +1,4 @@
+libopenmpi-dev
+openmpi-bin
+
+

train_settings/dvd/improved_diffusion/dist_util.py

@@ -9,14 +9,14 @@ import socket
 import blobfile as bf
 import torch as th
 import torch.distributed as dist
-try:
-    from mpi4py import MPI
-except ImportError:
-    MPI = None
-# from mpi4py import MPI
-
-if MPI is None or not hasattr(MPI, "COMM_WORLD"):
-    print("MPI not available, running in single process mode.")
+# try:
+#     from mpi4py import MPI
+# except ImportError:
+#     MPI = None
+from mpi4py import MPI
+
+# if MPI is None or not hasattr(MPI, "COMM_WORLD"):
+#     print("MPI not available, running in single process mode.")
 
 # Change this to reflect your cluster layout.
 # The GPU for a given rank is (rank % GPUS_PER_NODE).