Jax MLP NN loss is nan, without "env PYTHONUNBUFFERED"

[YoungSeokPark]

My jax neural network get nan-loss even for first iteration without env variable "PYTHONUNBUFFERED=1"

and with "env PYTHONUNBUFFERED" it works fine. It is just a simple FCN with params less than 10K,

Is there anyone have any clues?

Any helps would be appreciates. Thanks.

[JiaYaobo]

would u share your source code?