Flax module has no attribute optim

[Mehi44]

I tried running this in the colab notebook :
!python visualhull.py --config configs/blender --data_dir data/nerf_synthetic/lego --voxel_dir data/voxel_dil7/lego --dilation 7 --thresh 1. --alpha_bkgd True

I recieved this error :

Traceback (most recent call last): File "visualhull.py", line 4, in <module> from nerf import utils File "/content/VaxNeRF/nerf/utils.py", line 35, in <module> class TrainState: File "/content/VaxNeRF/nerf/utils.py", line 36, in TrainState optimizer: flax.optim.Optimizer AttributeError: module 'flax' has no attribute 'optim'