fixing dm-haiku jax.xla error message

Author: sokrypton

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",