fix Lightning-AI#2407

ibeltagy · Jul 28, 2020 · 1115d6b · 1115d6b
1 parent f8a176c
commit 1115d6b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_lightning/trainer/distrib_data_parallel.py b/pytorch_lightning/trainer/distrib_data_parallel.py
@@ -528,7 +528,7 @@ def ddp_train(self, process_idx, q, model, is_master=False, proc_offset=0):
             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)