google-research / jax3d

Apache License 2.0
735 stars 96 forks source link

Problem about flax.optim in MobileNerf #116

Closed AuthorityWang closed 2 years ago

AuthorityWang commented 2 years ago

When I try to train the mobilenerf model, I find a problem about flax. When I following the guide to install all the lib and trying to run stage1.py, it will give an error module 'flax' has no attribute 'optim'for line 1476 state=flax.optim.Adam(**adam_kwargs).create(model_vars). I find out that flax I installed is version0.6.0, but an older version like 0.5.3 will be ok. So maybe the requirements should make a change? Thanks :)

czq142857 commented 2 years ago

Thanks for bringing this up! Yes, optim has been removed in the latest flax package. Please use flax<0.6.0.