Adds torch.cuda.set_device calls to DDP examples (#1142)

Add set_device calls to DDP examples

Author: subramen

distributed/ddp-tutorial-series/multigpu.py

@@ -19,6 +19,7 @@ def ddp_setup(rank, world_size):
     os.environ["MASTER_ADDR"] = "localhost"
     os.environ["MASTER_PORT"] = "12355"
     init_process_group(backend="nccl", rank=rank, world_size=world_size)
+    torch.cuda.set_device(rank)
 
 class Trainer:
     def __init__(

distributed/ddp-tutorial-series/multigpu_torchrun.py

@@ -12,6 +12,7 @@
 
 def ddp_setup():
     init_process_group(backend="nccl")
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
 
 class Trainer:
     def __init__(

distributed/ddp-tutorial-series/multinode.py

@@ -12,6 +12,7 @@
 
 def ddp_setup():
     init_process_group(backend="nccl")
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
 
 class Trainer:
     def __init__(

distributed/minGPT-ddp/mingpt/main.py

@@ -8,6 +8,7 @@
 
 def ddp_setup():
     init_process_group(backend="nccl")
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
 
 def get_train_objs(gpt_cfg: GPTConfig, opt_cfg: OptimizerConfig, data_cfg: DataConfig):
     dataset = CharDataset(data_cfg)