Error Importing Flax Module

[Protinon]

I've followed all the steps outlined in the readme, including jaxlib and pytorch. This is python version 3.9.6

(venv) PS D:\DallE> .\venv\Scripts\python.exe .\dalle-playground\backend\app.py
--> Starting DALL-E Server. This might take up to two minutes.
2022-06-12 23:52:23.958841: W external/org_tensorflow/tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'cudart64_110.dll'; dlerror: cudart64_110.dll not found
2022-06-12 23:52:23.974713: W external/org_tensorflow/tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'cudart64_110.dll'; dlerror: cudart64_110.dll not found
Traceback (most recent call last):
  File "D:\DallE\dalle-playground\backend\app.py", line 13, in <module>
    from dalle_model import DalleModel
  File "D:\DallE\dalle-playground\backend\dalle_model.py", line 11, in <module>
    from dalle_mini import DalleBart, DalleBartProcessor
  File "D:\DallE\venv\lib\site-packages\dalle_mini\__init__.py", line 3, in <module>
    from .model import DalleBart, DalleBartProcessor
  File "D:\DallE\venv\lib\site-packages\dalle_mini\model\__init__.py", line 2, in <module>
    from .modeling import DalleBart
  File "D:\DallE\venv\lib\site-packages\dalle_mini\model\modeling.py", line 28, in <module>
    from flax.linen import partitioning as nn_partitioning
ImportError: cannot import name 'partitioning' from 'flax.linen' (D:\DallE\venv\lib\site-packages\flax\linen\__init__.py)

[michael-wise]

The requirements.txt file in ./backend does not specify a flax version. A newer version, such as below, includes flax.linen.partitioning

pip install flax==0.5.1

(don't forget to do this inside of your venv)

[Protinon]

Thanks! That got rid of the error.