Checkpoint on NaN

[agrimgupta92]

I am trying to debug a NaN error which happens during the middle of training i.e. after some epochs. After I detect the NaN I would like to save the model weights to inspect later. I was unable to find a way to do the same. Any help in this regard would be appreciated. Note I am using DDP so can't directly save the state dict.