fix(inference): update flax version (temporary fix)

dev/inference/inference_pipeline.ipynb

@@ -6,7 +6,7 @@
       "name": "DALL·E mini - Inference pipeline.ipynb",
       "provenance": [],
       "collapsed_sections": [],
-      "authorship_tag": "ABX9TyMUjEt1XMLq+6/GhSnVFsSx",
+      "authorship_tag": "ABX9TyOvSyRBA5cMjtEixUxYuduO",
       "include_colab_link": true
     },
     "kernelspec": {
@@ -826,7 +826,7 @@
         "id": "uzjAM2GBYpZX"
       },
       "source": [
-        "!pip install -q transformers flax\n",
+        "!pip install -q transformers flax==0.3.5\n",
         "!pip install -q git+https://github.com/patil-suraj/vqgan-jax.git  # VQGAN model in JAX\n",
         "!pip install -q git+https://github.com/borisdayma/dalle-mini.git  # Model files"
       ],