pushing A100 fix (default jaxlib not compatible with A100)

AlphaFold2.ipynb

@@ -148,7 +148,8 @@
         "  print(\"installing colabfold...\")\n",
         "  os.system(\"pip install -q --no-warn-conflicts 'colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/ColabFold'\")\n",
         "  if os.environ.get('TPU_NAME', False) != False:\n",
-        "    os.system(\"pip install -q --no-warn-conflicts -U dm-haiku==0.0.10 jax==0.3.25\")\n",
+        "    os.system(\"pip uninstall -y jax jaxlib\")\n",
+        "    os.system(\"pip install --no-warn-conflicts --upgrade dm-haiku==0.0.10 'jax[cuda12_pip]'==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\")\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",
         "  os.system(\"touch COLABFOLD_READY\")\n",
@@ -635,4 +636,4 @@
       ]
     }
   ]
-}
+}