I just came back to update this post and saw this reply, which is incidentally very close to what I have been doing.
My plan was to build in protecting in the model against the nans by saving the model_state_dict after each epoch and then if nans are detected in an epoch I would just reload the previous epochs model, lower the learning rate a bit and continue training.
However that doesn’t seem to work, because whenever I start to train with a model loaded like this the model parameters all goes to nan after using optimizer.step()
I don’t see how this can happen since the loss itself looks good, and I checked all the model weights before loss.backwards() and they also all looks good and reasonable, and I checked all the gradients of the weigths after loss.backwards(), but before optimizer.step() and they also all look good and reasonable, yet despite this everything goes nan.
When I reload the model I do not reload the optimizer (using adam), maybe that is the reason why this doesn’t work, though I can’t really see why this would be a problem.
Here is the method I use to check up on the model, this is used just after loss.backwards() and just before optimizer.step()
        weights = optimizer.param_groups[0]['params']
        weights_flat = [torch.flatten(weight) for weight in weights]
        weights_1d = torch.cat(weights_flat)
        assert not torch.isnan(weights_1d).any()
        assert not torch.isinf(weights_1d).any()
        print(f"{weights_1d.max()}, {weights_1d.min()}")

        grad_flat = [torch.flatten(weight.grad) for weight in weights]
        grad_1d = torch.cat(grad_flat)
        assert not torch.isnan(grad_1d).any()
        assert not torch.isinf(grad_1d).any()
        print(f"{grad_1d.max()}, {grad_1d.min()}")

Based on your suggestion I tried loading the model again, and enabling anomaly_detection, however strangely enough it didn’t seem to do anything, once again the model was loaded, data was passed through the model, a loss was generated and backwards propagated.
For the model weights / gradients - max / min I got:
8.083415985107422, -9.480425834655762
0.3138473927974701, -0.9406746625900269

Then I let the optimizer take a step and afterwards all the weights were nan again, but the anomaly detection didn’t seem to catch anything, strangely enough.
I will try to also save and load the optimizer in my current approach, but any other suggestions for how to debug this would be greatly appreciated.