Temporary fix for jax update in colab

Author: milot-mirdita

AlphaFold2.ipynb

@@ -148,6 +148,8 @@
         "  if os.environ.get('TPU_NAME', False) != False:\n",
         "    os.system(\"pip uninstall -y jax jaxlib\")\n",
         "    os.system(\"pip install --no-warn-conflicts --upgrade dm-haiku==0.0.10 'jax[cuda12_pip]'==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\")\n",
+        "  else:\n",
+        "    os.system(\"pip install --no-warn-conflicts --upgrade dm-haiku==0.0.13\")\n",
         "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabfold colabfold\")\n",
         "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/alphafold alphafold\")\n",
         "  os.system(\"touch COLABFOLD_READY\")\n",