[jcims]

You need to install flax 0.4.2. If you're using collab you just open a terminal (icon in the bottom left of the screen) and run:

    pip3 install flax==0.4.2

[kuprel]

Yes this is the fix for now. I need to address what is actually causing the dtype mismatch