search

google / orbax

Orbax provides common checkpointing and persistence utilities for JAX users
https://orbax.readthedocs.io/
Apache License 2.0
307 stars 36 forks source link

Can't save model on Colab but can on Kaggle #1229

Closed windmaple closed 1 month ago

windmaple commented 1 month ago

I have a very simple model and orbax won't let me save it: https://drive.google.com/file/d/1hM9NN5vY2EO_NvieIFZ_A4uHUG-s5VEE/view?usp=sharing

Same code works fine on Kaggle though: https://www.kaggle.com/code/windmaple/notebook28117e6799

windmaple commented 1 month ago

Never mind. I had to install orbax-checkpoint instead of orbax.