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]>

Author: ibeltagy

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)