I'm trying to run RegNeRF, which requires flax. On installing the latest version of flax==0.6.0, I got an error stating flax has no attribute optim. This answer suggested to downgrade flax to 0.5.1. On doing that, now I'm getting the error AttributeError: module 'flax' has no attribute 'nn'
I could not find any solutions on the web for this error. Any help is appreciated.
I'm using ubuntu 20.04