diff --git a/batch/AlphaFold2_batch.ipynb b/batch/AlphaFold2_batch.ipynb
index 6dd66cb7..910d3858 100644
--- a/batch/AlphaFold2_batch.ipynb
+++ b/batch/AlphaFold2_batch.ipynb
@@ -127,9 +127,11 @@
         "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\" \"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",
+        "  pip install -q --no-warn-conflicts \"colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/ColabFold\"\n",
+        "  pip install --upgrade dm-haiku\n",
+        "  ln -s /usr/local/lib/python3.*/dist-packages/colabfold colabfold\n",
+        "  ln -s /usr/local/lib/python3.*/dist-packages/alphafold alphafold\n",
+        "  sed -i 's/weights = jax.nn.softmax(logits)/logits=jnp.clip(logits,-1e8,1e8);weights=jax.nn.softmax(logits)/g' alphafold/model/modules.py\n",
         "  touch COLABFOLD_READY\n",
         "fi\n",
         "\n",