Using mix precision, the loss flattens out after the first few iterations. The model trains fine when mix precision is not used.
Here is an example of how it’s implemented. I am using a ctc loss function
for i, _data in enumerate(train_loader):
    spectrograms, labels, input_lengths, label_lengths = _data 
    spectrograms, labels = spectrograms.cuda(), labels.cuda()
    with autocast():
        output = model(spectrograms)  # (batch, time, n_class)
        output = F.log_softmax(output, dim=2)

        loss = ctc_loss(output, labels, input_lengths, label_lengths)
        loss = loss / args.grad_acc_steps
        scaler.scale(loss).backward()

    if i % args.grad_acc_steps == 0:
        scaler.unscale_(optimizer)  # unscale to clip gradient
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)
        scaler.update()
        scheduler.step()