Hi,
what GPU are you running?
Also Flax version 1.10 is pretty old now and not really supported anymore.
You might want to try Flax 1.12 and see if you encounter the same issue there. There have been a lot of fixes in 1.11 and 1.12.
