Could you rerun the code with torch.autograd.set_detect_anomaly(True) at the beginning of the script and post the stack trace here?
Based on the description I assume the loss does not contain any invalid values?
If that’s the case, could you check all gradients after the first backward() pass for NaN values?