Explicitly set ddp device

cxmlg · Feb 4, 2023 · 9da1627 · 9da1627
1 parent 3fd4c0c
commit 9da1627
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: