AttributeError: module 'flax' has no attribute 'optim'

Question

My code is as follows:

!pip install flax
init_params = TransporterNets().init(key, init_img, init_text, init_pix)['params']
print(f'Model parameters: {n_params(init_params):,}')
optim = flax.optim.Adam(lr=1e-4).create(init_params)

However it shows the following error:

AttributeError: module 'flax' has no attribute 'optim'

Even though I have seen documentation of optim attribute in flax module. How to fix it?

did you import it first ? `import flax ` – mrCopiCat Aug 25 '22 at 14:04 — mrCopiCat, Aug 25 '22 at 14:04
@mrCopiCat Yes, I have imported it. Still doesn't work. – Md Tawsif Mostafiz 170021031 Aug 25 '22 at 14:29 — Md Tawsif Mostafiz 170021031, Aug 25 '22 at 14:29

score 2 · Accepted Answer · answered Aug 26 '22 at 12:35

2

You can temporarily solve the issue by downgrading flax version from 0.6.0 to 0.5.1 at the moment.

pip install flax==0.5.1

answered Aug 26 '22 at 12:35

Baris CELIK

36
1

AttributeError: module 'flax' has no attribute 'optim'

1 Answers1

Linked