Using Lorentz model returns nan about differentiable-frechet-mean HOT 2 OPEN

Hi. It could be for a variety of reasons, depending on your situation. Could you share your test code?

from differentiable-frechet-mean.

I used the demo code for batch normalization provided by you but using Lorentz model instead of Poincare. It turned out to be a machine precision problem. When I use float32 tensors, it returns nan but with float64 it works. However, this makes the model twice as large and much more slower to learn. Any ideas on how to modify the code to make it works with float32 tensors?

from differentiable-frechet-mean.