Merge pull request karpathy#116 from ramtingh/master

Minor change to allow using ddp with exclusive process mode
kustomzone · Feb 4, 2023 · 53d56b8 · 53d56b8
2 parents 3fd4c0c + 9da1627
commit 53d56b8
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/train.py b/train.py
@@ -85,6 +85,7 @@
     ddp_rank = int(os.environ['RANK'])
     ddp_local_rank = int(os.environ['LOCAL_RANK'])
     device = f'cuda:{ddp_local_rank}'
+    torch.cuda.set_device(device)
     master_process = ddp_rank == 0 # this process will do logging, checkpointing etc.
     seed_offset = ddp_rank # each process gets a different seed
 else: