Nan while training

[rusheniii]

Is your feature request related to a problem? Please describe. When you are training an LDDMM atlas, if you don't set the parameters correctly, the loss can go to Nan.

Describe the solution you'd like When the training loss is Nan, you should stop training. Currently, the code continues training for X number of iterations.