pin jax/tensorflow versions

sokrypton · Feb 28, 2023 · 60e2ee0 · 60e2ee0
1 parent d0bd141
commit 60e2ee0
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/AlphaFold2.ipynb b/AlphaFold2.ipynb
@@ -144,9 +144,9 @@
         "  echo \"installing colabfold...\"\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.5.1\"\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.5.1\" \"tensorflow-cpu==2.7.4\"\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",
         "\n",
         "  # for debugging\n",
         "  ln -s /usr/local/lib/python3.*/dist-packages/colabfold colabfold\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.5.1\"\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.5.1\" \"tensorflow-cpu==2.7.4\"\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",