There could be many reasons why this could be the case,
I will suggest to turn on the flag torch.autograd.set_detect_anomaly(True)
This will drop in error wherever the autograd engine encounters NaN and give you the function where the error occurred, It could be advantageous to debug and see what is going on!