Issue running models

[minhng22]

Steps to reproduce:

Run python transformer/ht_main.py --alsologtostderr \ --gin_file=base_htrans.gin \ --gin_file=size/small_test.gin

WARNING:absl:GlobalAsyncCheckpointManager is not imported correctly. Checkpointing of GlobalDeviceArrays will not be available.To use the feature, install tensorstore. Traceback (most recent call last): File "/Users/minhnguyen/Documents/meliad/transformer/ht_main.py", line 25, in <module> from transformer import launcher File "/Users/minhnguyen/Documents/meliad/transformer/launcher.py", line 23, in <module> import training_loop File "/Users/minhnguyen/Documents/meliad/transformer/training_loop.py", line 32, in <module> import optimizer_config as opt_config File "/Users/minhnguyen/Documents/meliad/transformer/optimizer_config.py", line 21, in <module> from flax import optim ImportError: cannot import name 'optim' from 'flax' (/Users/minhnguyen/opt/anaconda3/lib/python3.9/site-packages/flax/__init__.py)

I have followed all steps to install in README. Env is MacOS.

[jmsalvador2395]

hi @minhng22. Just ran into the same problem. I found this link that basically says that the 'optim' module was removed from flax starting from version 0.6.0.

I was able to downgrade flax by running
pip install --upgrade flax==0.5.3

hope this helps

[minhng22]

Yes confirmed that this fixed the issue. Thanks,