Error on Nerfies Training v2.ipynb

[Juan Alonso]

Hi!

I'm trying to run the notebook Nerfies Training v2.ipynb (using the Colab environment) and I'm getting this error in the "Configuration" cell, both for GPU and TPU.

ImportError                               Traceback (most recent call last)

<ipython-input-5-be45ab41e1a4> in <module>()
      6 from IPython.display import display, Markdown
      7 
----> 8 from nerfies import configs
      9 
     10 

/usr/local/lib/python3.7/dist-packages/nerfies/configs.py in <module>()
     17 
     18 import dataclasses
---> 19 from flax import nn
     20 import gin
     21 import immutabledict

ImportError: cannot import name 'nn' from 'flax' (/usr/local/lib/python3.7/dist-packages/flax/__init__.py)

It seems that flax.nn is deprecated and flax.linen should be used instead.

Thank you very much

[zach]

I replaced line 18 with from flax import linen as nn and then the notebook ran as expected

I just submitted a pull request to fix this