I just updated the jax to 0.4.14 on the server of our lab. But after training by the same file that had been run in jax 0.3.25, I found that the checkpoint file cannot be loaded by my test file, no matter in 0.4.14 or 0.3.25. Moreover, the save path changed to "./filename/checkpoint" rather than "./filename" in 0.3.25, although the code is the same. I wonder if anyone meets with the same problem and how I should fix it. Should I retreat to 0.3.25? The errors are below.problem while loading
I tried updating my environment by yaml but it still didn't work. The errors were the same.