Let’s say I have limited insight into what exactly is happening. (For this it would be interesting to see 1) the value of the param & optimizer state just before it turns NaN 2) the value of the gradient that causes the state to turn NaN.
So the epsilon helps when v_hat is 0, i.e. when v is 0.



 Azerus:

So I tried this suggestion, but with eps=1e-4 my loss hardly decrease over time while the loss was decreasing quite fast when using the default eps=1e-8


This seems to tell me that the loss gradient is terribly small all along.
One other thing you might try is to switch Adam, which works at a coefficient level, for an optimizer that works at a tensor level like LAMB (there are more, and probably one that is closer to Adam in spirit, but this is the acronym I can remember).
Best regards
Thomas