Thanks for the update. Could you check for NaN values in all .grad attributes before calling optimizer.step()?
If the input and output are fine, but the step() call creates the NaNs, it seems that the gradients must become invalid values.
You can check it via:
for name, param in model.named_parameters():
    print(name, torch.isfinite(param.grad).all())

after the backward call.
If you get an invalid value in the gradients, could you store the input, target as well as the model.state_dict and upload it somewhere so that we could try to reproduce and debug this issue?