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

Lightning-AI · Apr 3, 2020 · 5fc8165 · 5fc8165
1 parent 868b172
commit 5fc8165
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
@@ -638,4 +638,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