I would recommend to add torch.autograd.set_detect_anomaly(True) at the beginning of your script, which would print a stack trace showing pointing towards the operation, which created the first NaN output.
This should be helpful in debugging the issue.
PS: you can post code snippets by wrapping them into three backticks ```, which would make debugging a bit easier.