diff --git a/AlphaFold2.ipynb b/AlphaFold2.ipynb
index b5d87666..a01e669c 100644
--- a/AlphaFold2.ipynb
+++ b/AlphaFold2.ipynb
@@ -154,6 +154,7 @@
         "if not os.path.isfile(\"COLABFOLD_READY\"):\n",
         "  print(\"installing colabfold...\")\n",
         "  os.system(\"pip install -q --no-warn-conflicts 'colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/ColabFold'\")\n",
+        "  os.system(\"pip install --upgrade dm-haiku\")\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",
         "  # patch for jax > 0.3.25\n",