Repository: Lightning-AI/pytorch-lightning

Issue:
Incorrect default cuda device when using single gpu other than cuda:0
## 🐛 Bug

The default `cuda` is not set properly to the `trainer.root_gpu` in single-GPU mode. The tensors created with `device='cuda'` will be placed on the incorrect gpu, and the dataloader will acquire memory on the incorrect gpu when `pin_memory=True`.

Maybe we'll need to add
`torch.cuda.set_device(self.trainer.root_gpu)` to https://github.com/PyTorchLightning/pytorch-lightning/blob/5dfc7b157e7febab692036b7392dac8b52f41b87/pytorch_lightning/accelerators/gpu_backend.py#L24
as `DDPBackend` did:

https://github.com/PyTorchLightning/pytorch-lightning/blob/5dfc7b157e7febab692036b7392dac8b52f41b87/pytorch_lightning/accelerators/ddp_backend.py#L195

### To Reproduce

Running the following code will get 

`RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0!`

#### Code sample

```python
import pytorch_lightning as pl
import torch
from torch import nn
from torch.utils import data


class Dataset(data.Dataset):

  def __getitem__(self, item):
    return torch.zeros(1)

  def __len__(self):
    return 5


class Model(pl.LightningModule):

  def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.x = nn.Parameter(torch.zeros(1))

  def forward(self, *args, **kwargs):
    return self.x

  def training_step(self, *args, **kwargs):
    return self.x + torch.zeros(1, device='cuda')  # RuntimeError.

  def train_dataloader(self):
    return data.DataLoader(Dataset(), num_workers=1, pin_memory=True)

  def configure_optimizers(self):
    return torch.optim.SGD(self.parameters(), 1.0)


if __name__ == '__main__':
  trainer = pl.Trainer(gpus=[1], num_sanity_val_steps=0, max_epochs=1)
  model = Model()
  trainer.fit(model)

```

### Expected behavior

No `RuntimeError` occurs.

### Environment

* CUDA:
	- GPU:
	- available:
	- version:
* Packages:
	- numpy:             1.18.5
	- pyTorch_debug:     False
	- pyTorch_version:   1.6.0
	- pytorch-lightning: 0.9.0rc16
	- tensorboard:       2.3.0
	- tqdm:              4.48.2
* System:
	- OS:                Windows
	- architecture:
		- 64bit
		- WindowsPE
	- processor:
	- python:            3.7.3
	- version:           10.0.18362

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Directly related to #3016 I believe, perhaps a duplicate? Thanks for bringing this up 😄 

Resolving patch:
diff --git a/pytorch_lightning/accelerators/gpu_backend.py b/pytorch_lightning/accelerators/gpu_backend.py
--- a/pytorch_lightning/accelerators/gpu_backend.py
+++ b/pytorch_lightning/accelerators/gpu_backend.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import torch
 from pytorch_lightning.core import LightningModule
 from pytorch_lightning.utilities import AMPType
 
@@ -32,6 +33,7 @@ def setup(self, model):
         # call setup
         self.trainer.call_setup_hook(model)
 
+        torch.cuda.set_device(self.trainer.root_gpu)
         model.cuda(self.trainer.root_gpu)
 
         # CHOOSE OPTIMIZER


Test patch: