"Hello, I am trying to run the 01_inference_pretrained.ipynb notebook on a standard Google Colab T4 GPU runtime. I have tried following the official installation script, as well as manually forcing dependencies like jax==0.4.20, jaxlib==0.4.20, and numpy<2, but I consistently run into contradictory dependency errors. The final error, even with the correct jax versions apparently installed, is AttributeError: module 'jax.tree_util' has no attribute 'register_dataclass'.I have spent hours but didn't figure this out.
Could you please provide an updated, working set of installation commands for the current Colab environment? Thank you."
"Hello, I am trying to run the 01_inference_pretrained.ipynb notebook on a standard Google Colab T4 GPU runtime. I have tried following the official installation script, as well as manually forcing dependencies like jax==0.4.20, jaxlib==0.4.20, and numpy<2, but I consistently run into contradictory dependency errors. The final error, even with the correct jax versions apparently installed, is AttributeError: module 'jax.tree_util' has no attribute 'register_dataclass'.I have spent hours but didn't figure this out.
Could you please provide an updated, working set of installation commands for the current Colab environment? Thank you."