Bi-level optimization and Adam causing NaNs

[averageFlaxUser]

I have a bi-level optimization where model A produces some output vector z that is used as a regularization in the loss_fn of model B that is optimized on some image data. Assuming my implementation is correct, I noticed that when model B is optimized via Adam, gradients become nan but any other optimizer works fine. After some hours of digging I noticed the eps_root hyper-parameter is set by default to 0.0. Changing this value fixes the issue. My concern is, why is this so. Is this an issue from my implementation of this is expected in some cases?