Then it is something else, probably. If you use sampling from trainable distributions, the issue can be there.
Generally autograd.set_detect_anomaly(True) should show the problematic part (NaN inducing). As it slows down training, it is better to enable it late.