Repository: Lightning-AI/pytorch-lightning

Issue:
DDP on GPUs invalid ordinal
<!-- 
### Common bugs:
1. Tensorboard not showing in Jupyter-notebook see [issue 79](https://github.com/PyTorchLightning/pytorch-lightning/issues/79).    
2. PyTorch 1.1.0 vs 1.2.0 support [see FAQ](https://github.com/PyTorchLightning/pytorch-lightning#faq)    
-->

## 🐛 Bug

On latest version (master), training with DDP backend on GPUs (let's say 6,7) results in a CUDA error "Invalid device ordinal." 

### To Reproduce

Steps to reproduce the behavior:

1. Run any Lightning Module with DDP backend on more than 1 GPU with GPU indexes that do not start from 0
2. See error

```
INFO:lightning:GPU available: True, used: True
INFO:lightning:VISIBLE GPUS: 6,7
WARNING:lightning:SLURM_NODEID or NODE_RANK environment variable is not defined. Set as 0.
WARNING:lightning:MASTER_ADDR environment variable is not defined. Set as localhost
WARNING:lightning:SLURM_NODEID or NODE_RANK environment variable is not defined. Set as 0.
WARNING:lightning:MASTER_ADDR environment variable is not defined. Set as localhost
THCudaCheck FAIL file=/pytorch/torch/csrc/cuda/Module.cpp line=59 error=101 : invalid device ordinal
THCudaCheck FAIL file=/pytorch/torch/csrc/cuda/Module.cpp line=59 error=101 : invalid device ordinal
Traceback (most recent call last):
  File "bin/run.py", line 110, in <module>
    run(config, args.mode)
  File "bin/run.py", line 86, in run
    trainer.fit(system)
  File "/home/kevin/miniconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 744, in fit
    mp.spawn(self.ddp_train, nprocs=self.num_processes, args=(model,))
  File "/home/kevin/miniconda3/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 171, in spawn
    while not spawn_context.join():
  File "/home/kevin/miniconda3/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 118, in join
    raise Exception(msg)
Exception:

-- Process 1 terminated with the following error:
Traceback (most recent call last):
  File "/home/kevin/miniconda3/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 19, in _wrap
    fn(i, *args)
  File "/home/kevin/miniconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/distrib_data_parallel.py", line 332, in ddp_train
    torch.cuda.set_device(self.root_gpu)
  File "/home/kevin/miniconda3/lib/python3.7/site-packages/torch/cuda/__init__.py", line 292, in set_device
    torch._C._cuda_setDevice(device)
RuntimeError: cuda runtime error (101) : invalid device ordinal at /pytorch/torch/csrc/cuda/Module.cpp:59

```

### Expected behavior

I expect this error to not occur.

### Environment

* CUDA:
        - GPU:
                - TITAN RTX
                - TITAN RTX
                - TITAN RTX
                - TITAN RTX
                - TITAN RTX
                - TITAN RTX
                - TITAN RTX
                - TITAN RTX
        - available:         True
        - version:           10.1
* Packages:
        - numpy:             1.17.4
        - pyTorch_debug:     False
        - pyTorch_version:   1.4.0
        - pytorch-lightning: 0.7.4rc1
        - tensorboard:       2.0.0
        - tqdm:              4.45.0
* System:
        - OS:                Linux
        - architecture:
                - 64bit
                - 
        - processor:         x86_64
        - python:            3.7.4
        - version:           #97-Ubuntu SMP Wed Apr 1 03:25:46 UTC 2020

### Additional context

My best guess is that due to the new refactoring for ddp cpu backend, setting the CUDA device to `root_gpu` instead of using the provided `gpu_idx` parameter in `ddp_train()` is causing the error. Specifically, for 2 gpu, it's using index 6 instead of index 0 and index 7 instead of index 1 since we're setting `CUDA_VISIBLE_DEVICES=6,7`.



Discussion:
Hi! thanks for your contribution!, great first issue!
good catch. mind submitting a PR?

Resolving patch:
diff --git a/pytorch_lightning/trainer/distrib_data_parallel.py b/pytorch_lightning/trainer/distrib_data_parallel.py
--- a/pytorch_lightning/trainer/distrib_data_parallel.py
+++ b/pytorch_lightning/trainer/distrib_data_parallel.py
@@ -327,7 +327,7 @@ def ddp_train(self, process_idx, model):
         # MODEL
         # copy model to each gpu
         if self.on_gpu:
-            self.root_gpu = self.data_parallel_device_ids[process_idx]
+            self.root_gpu = process_idx
             torch.cuda.set_device(self.root_gpu)
             model.cuda(self.root_gpu)
 


Test patch: