diff --git a/alignn/train.py b/alignn/train.py
index aedfd61..ddfe602 100644
--- a/alignn/train.py
+++ b/alignn/train.py
@@ -36,6 +36,7 @@
 
 
 def setup(rank, world_size):
+    """Set up multi GPU rank."""
     os.environ["MASTER_ADDR"] = "localhost"
     os.environ["MASTER_PORT"] = "12355"
     # Initialize the distributed environment.
diff --git a/alignn/train_alignn.py b/alignn/train_alignn.py
index 9aaf4e5..cdc61a1 100644
--- a/alignn/train_alignn.py
+++ b/alignn/train_alignn.py
@@ -396,7 +396,7 @@ def train_for_folder(
 if __name__ == "__main__":
     args = parser.parse_args(sys.argv[1:])
     world_size = int(torch.cuda.device_count())
-    rank = [0, 1]
+    print("world_size", world_size)
     torch.multiprocessing.spawn(
         train_for_folder,
         args=(