Lint fix.

usnistgov · Apr 28, 2024 · 6e19d56 · 6e19d56
1 parent 63ba5e3
commit 6e19d56
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git 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
@@ -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=(