Sure, in case you want to debug a bit further:

Add torch.autograd.set_detect_anomaly(True) at the beginning of your script. This would yield a stack trace with the operation, which caused the first NaN output.
If you are using mixed-precision training (via native amp, apex, or your manual implementation), disable it for the sake of debugging.
Update to the latest PyTorch version, if not already done.
Check all torch.log, torch.sqrt, torch.pow calls etc. which might easily create Infs/NaNs.

Let me know, if you can release the model or a smaller fake example, which would reproduce this issue.