diff --git a/dlio_benchmark/main.py b/dlio_benchmark/main.py
index 0d56bb20..7d7b6fb7 100644
--- a/dlio_benchmark/main.py
+++ b/dlio_benchmark/main.py
@@ -16,7 +16,8 @@
 """
 import os
 import math
-
+from mpi4py import MPI
+comm = MPI.COMM_WORLD
 import logging
 from time import time, sleep
 import json
@@ -49,9 +50,8 @@
 from dlio_benchmark.utils.utility import Profile, PerfTrace
 
 dlp = Profile(MODULE_DLIO_BENCHMARK)
-from mpi4py import MPI
 # To make sure the output folder is the same in all the nodes. We have to do this. 
-MPI.COMM_WORLD.Barrier()
+comm.Barrier()
 import hydra
 
 class DLIOBenchmark(object):