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

shubhamagarwal92 · Mar 15, 2020 · cc34c94 · cc34c94
1 parent cdd10f9
commit cc34c94
Showing 1 changed file with 0 additions and 1 deletion.
diff --git a/pytorch_lightning/trainer/distrib_parts.py b/pytorch_lightning/trainer/distrib_parts.py
@@ -643,7 +643,6 @@ def determine_root_gpu_device(gpus):
     # 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)
     root_device = (torch.device("cuda", root_gpu) if root_gpu >= 0 else torch.device("cpu"))
     torch.cuda.set_device(root_device)