Hi @bshabash, glad to hear you were able to fix. I’d be happy to help you add a checkpoint or validation logic if you can formulate a good idea of where the NAN error arose. In general I find NAN errors are difficult to track; we have some NAN checking in Pyro but it is difficult to trace NANs through PyTorch gradients, especially as NAN propagation rules are subtle and sometimes change across PyTorch releases.