SA: for Lightning-AI#958: set torch cuda device when finding root

shubhamagarwal92 · Mar 12, 2020 · 8a66308 · 8a66308
1 parent 5e013f6
commit 8a66308
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/pytorch_lightning/trainer/distrib_parts.py b/pytorch_lightning/trainer/distrib_parts.py
@@ -640,4 +640,10 @@ def determine_root_gpu_device(gpus):
     # set root gpu
     root_gpu = gpus[0]
 
+    # set cuda device to root gpu
+    # related to https://github.com/PyTorchLightning/pytorch-lightning/issues/958
+    # Refer solution: https://github.com/pytorch/pytorch/issues/9871#issuecomment-408304190
+    root_device = torch.device("cuda", root_gpu)
+    torch.cuda.set_device(root_device)
+
     return root_gpu