Thank you for the advice. I’ve added the gradient clipping as you suggested, but the loss is still nan. The value in args.clip_grad is really large though, so I don’t think it is doing anything, either way, just a simple way to catch huge gradients. But I agree it should be in both conditions for consistency.
        if nb_tr_steps % args.grad_acum_steps == 0:
            if args.use_apex:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip_grad)
                scaler.step(optimizer)
                scaler.update()
            else:
                torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip_grad)
                optimizer.step()
            optimizer.zero_grad()

Is there anything else I could try?