Could you check, if all inputs contain valid values via torch.isfinite(input).all() and torch.isfinite(output).all()?
If the data is alright, I would recommend to set torch.autograd.detect_anomaly(True) at the beginning of the script and post the stack trace here. This should hopefully point to the operation, which created the NaN outputs.