pin jax version

sokrypton · Feb 28, 2023 · 776071a · 776071a
1 parent 6b1cf7b
commit 776071a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/AlphaFold2.ipynb b/AlphaFold2.ipynb
@@ -212,9 +212,9 @@
         "if [ ! -f COLABFOLD_READY ]; then\n",
         "  # install dependencies\n",
         "  # We have to use \"--no-warn-conflicts\" because colab already has a lot preinstalled with requirements different to ours\n",
-        "  pip install -q --no-warn-conflicts \"colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/ColabFold@v1.4.0\"\n",
-        "  # high risk high gain\n",
-        "  pip install -q \"jax[cuda11_cudnn805]>=0.3.8,<0.4\" -f https://storage.googleapis.com/jax-releases/jax_releases.html\n",
+        "  pip install -q --no-warn-conflicts \"colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/ColabFold@v1.4.0\" \"tensorflow-cpu==2.11.0\"\n",
+        "  pip uninstall -yq jax jaxlib\n",
+        "  pip install -q \"jax[cuda]==0.3.25\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
         "  touch COLABFOLD_READY\n",
         "fi\n",
         "\n",

diff --git a/batch/AlphaFold2_batch.ipynb b/batch/AlphaFold2_batch.ipynb
@@ -119,9 +119,9 @@
         "if [ ! -f COLABFOLD_READY ]; then\n",
         "  # install dependencies\n",
         "  # We have to use \"--no-warn-conflicts\" because colab already has a lot preinstalled with requirements different to ours\n",
-        "  pip install -q --no-warn-conflicts \"colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/ColabFold@v1.4.0\"\n",
-        "  # high risk high gain\n",
-        "  pip install -q \"jax[cuda11_cudnn805]>=0.3.8,<0.4\" -f https://storage.googleapis.com/jax-releases/jax_releases.html\n",
+        "  pip install -q --no-warn-conflicts \"colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/ColabFold@v1.4.0\" \"tensorflow-cpu==2.11.0\"\n",
+        "  pip uninstall -yq jax jaxlib\n",
+        "  pip install -q \"jax[cuda]==0.3.25\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
         "  touch COLABFOLD_READY\n",
         "fi\n",
         "\n",