`precision=16` displaying wrong loss in progress bar
🐛 Bug
When training on a GPU (using native AMP) and setting precision=16, the loss displayed by the progress bar is some crazy large number.
Stopping the example bellow in the middle of Epoch 1 gives a loss of ~15 000. If I train with precision=32, this loss is the true value of ~0.23.
The loss tensor is OK, if I add a print statement in the training loop it displays normal values.
Code sample

import torch
from torch.nn import functional as F
from torch import nn

from pytorch_lightning.core.lightning import LightningModule
from pytorch_lightning import Trainer

from torch.utils.data import DataLoader, random_split
from torchvision.datasets import MNIST

import os
from torchvision import datasets, transforms

class LitMNIST(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer_1 = torch.nn.Linear(28 * 28, 128)
        self.layer_2 = torch.nn.Linear(128, 256)
        self.layer_3 = torch.nn.Linear(256, 10)

    def forward(self, x):
        batch_size, channels, width, height = x.size()
        x = x.view(batch_size, -1)
        x = self.layer_1(x)
        x = torch.relu(x)
        x = self.layer_2(x)
        x = torch.relu(x)
        x = self.layer_3(x)
        x = torch.log_softmax(x, dim=1)
        return x

    def train_dataloader(self):
        transform=transforms.Compose([transforms.ToTensor(),
                                      transforms.Normalize((0.1307,), (0.3081,))])
        mnist_train = MNIST(os.getcwd(), train=True, download=False, transform=transform)
        return DataLoader(mnist_train, batch_size=64)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)

        # add logging
        logs = {'loss': loss}
        return {'loss': loss, 'log': logs}
    

model = LitMNIST()
trainer = Trainer(gpus=1, precision=16)
trainer.fit(model)
Environment

PyTorch Version:  v1.7.0.dev20200704 (nightly)
OS (e.g., Linux): Ubuntu 18.04
How you installed PyTorch (conda, pip, source): conda
Python version: 3.8
CUDA/cuDNN version: 10.2 (installed with conda from pytorch chanel)
GPU models and configuration: RTX 2070 SUPER