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 
  
    
      pytorch-lightning/pytorch_lightning/accelerators/gpu_backend.py
    
    
         Line 24
      in
      5dfc7b1
    
  
  
    

        
          
           class GPUBackend(object): 
        
    
  


as DDPBackend did:

  
    
      pytorch-lightning/pytorch_lightning/accelerators/ddp_backend.py
    
    
         Line 195
      in
      5dfc7b1
    
  
  
    

        
          
           torch.cuda.set_device(self.trainer.root_gpu) 
        
    
  


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