fix selecting GPUs using CUDA_VISIBLE_DEVICES (#2739)

* fix #2407 * Update pytorch_lightning/trainer/distrib_data_parallel.py Co-authored-by: Jirka Borovec <[email protected]> Co-authored-by: Jirka Borovec <[email protected]>
Lightning-AI · Aug 2, 2020 · 38fce2e · 38fce2e
1 parent 7da7d2e
commit 38fce2e
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/pytorch_lightning/trainer/distrib_data_parallel.py b/pytorch_lightning/trainer/distrib_data_parallel.py
@@ -555,8 +555,7 @@ def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0
             gpu_idx = process_idx
             if is_master:
                 # source of truth is cuda for gpu idx
-                gpus = os.environ['CUDA_VISIBLE_DEVICES'].split(',')
-                gpu_idx = int(gpus[self.local_rank])
+                gpu_idx = self.local_rank
 
             self.root_gpu = gpu_idx
             torch.cuda.set_device(self.root_gpu)