Model and Input not on same GPU when training with native AMP and DP
🐛 Bug
Training with distributed_backend='dp' and precision=16 result in the error that some input/output of the model is not on the same GPU.
To Reproduce
from pytorch_lightning import Trainer, seed_everything
from pl_examples.models.lightning_template import LightningTemplateModel
seed_everything(234)

def main():
    # model = LightningTemplateModel(**vars(args))
    model = LightningTemplateModel()

    trainer = Trainer(
        gpus=2,
        num_nodes=1,
        distributed_backend='dp',
        precision=16
        )

    trainer.fit(model)

if __name__ == '__main__':
    main()

Run the above example code will result in following error:
Traceback (most recent call last):
  File "./_debug_pytorch_lightning_16gpu.py", line 28, in <module>
    main()
  File "./_debug_pytorch_lightning_16gpu.py", line 24, in main
    trainer.fit(model)
  File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 973, in fit
    self.dp_train(model)
  File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/trainer/distrib_parts.py", line 267, in dp_train
    self.run_pretrain_routine(model)
  File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1139, in run_pretrain_routine
    False)
  File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 291, in _evaluate
    output = self.evaluation_forward(model, batch, batch_idx, dataloader_idx, test_mode)
  File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 441, in evaluation_forward
    output = model(*args)
  File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/overrides/data_parallel.py", line 65, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/overrides/data_parallel.py", line 69, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/overrides/data_parallel.py", line 209, in parallel_apply
    raise output
  File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/overrides/data_parallel.py", line 172, in _worker
    output = module.validation_step(*input, **kwargs)
  File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pl_examples/models/lightning_template.py", line 102, in validation_step
    val_loss = F.cross_entropy(y_hat, y)
  File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/torch/nn/functional.py", line 2422, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/torch/nn/functional.py", line 2218, in nll_loss
    ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
RuntimeError: Assertion `THCTensor_(checkGPU)(state, 4, input, target, output, total_weight)' failed. Some of weight/gradient/input tensors are located on different GPUs. Please move them to a single one. at /opt/conda/conda-bld/pytorch_1593673617738/work/aten/src/THCUNN/generic/ClassNLLCriterion.cu:31

Expected behavior
The model should train successfully on 2 GPUs with native PyTorch AMP.
Environment

CUDA:
- GPU:
- Tesla V100-SXM2-16GB
- Tesla V100-SXM2-16GB
- Tesla V100-SXM2-16GB
- Tesla V100-SXM2-16GB
- available:         True
- version:           10.2
Packages:
- numpy:             1.18.5
- pyTorch_debug:     False
- pyTorch_version:   1.7.0.dev20200702
- pytorch-lightning: 0.8.5-dev
- tensorboard:       2.2.2
- tqdm:              4.47.0
System:
- OS:                Linux
- architecture:
- 64bit
-
- processor:         x86_64
- python:            3.7.7
- version:           #30~18.04.1-Ubuntu SMP Mon Jun 22 15:48:21 UTC 2020

Additional context
This only seems to be an issue with native AMP, not Apex AMP.