Loss not on correct device when training with move_metrics_to_cpu=True

[mibaumgartner]

🐛 Bug

move_metrics_to_cpu=True seems to move loss to CPU which results in an error when training with native mixed precision.
This is related to the original issue reported in MIC-DKFZ/nnDetection#25 and did not occur with other lightning versions (older version and/or move_metrics_to_cpu=False work fine).

Error when training with mixed precision and move_metrics_to_cpu=True:

/usr/local/lib/python3.7/dist-packages/torch/cuda/amp/grad_scaler.py in scale(self, outputs)
    159         # Short-circuit for the common case.
    160         if isinstance(outputs, torch.Tensor):
--> 161             assert outputs.is_cuda or outputs.device.type == 'xla'
    162             if self._scale is None:
    163                 self._lazy_init_scale_growth_tracker(outputs.device)

To Reproduce

Can be reproduced with the boring model in colab by passing the following flags to the trainer:

        precision=16,  # native mixed precision
        move_metrics_to_cpu=True,
        gpus=[0],  # use GPU

Expected behavior

No error :)

Environment

• PyTorch Lightning Version (e.g., 1.3.0):
• PyTorch Version (e.g., 1.8)
• Python version:
• OS (e.g., Linux):
• CUDA/cuDNN version:
• GPU models and configuration:
• How you installed PyTorch (conda, pip, source):
• If compiling from source, the output of torch.__config__.show():
• Any other relevant information:

Additional context

[tchaton]

Hey @mibaumgartner,

Would it be possible for you to share a reproducible script ?

Best,
T.C

[mibaumgartner]

Hi @tchaton,

here is the Boring Model Colab Notebook (latest lightning pypi release 1.4.5):
https://colab.research.google.com/drive/1pvHtF7Zor2OOLjE6NFjxxHTAassqdidb?usp=sharing

Best,
Michael

[tchaton]

Thanks @mibaumgartner,

I can reproduce the bug locally.

Best,
T.C