I assume “after the first batch” means that the first output and loss tensors are valid, while the second iteration produces a NaN output?
If that’s the case, could you check all gradients in the model using:
for name, param in model.named_parameters():
    print(name, torch.isfinite(param.grad).all())

after the first backward call and check, if some values are NaNs?
Alternatively, you could also set torch.autograd.set_detect_anomaly(True) at the beginning of the script, which should give you a stack trace pointing to the method, which created the NaNs in the backward pass.