fix: jax has no extras

borisdayma · Aug 22, 2023 · fd818d2 · fd818d2
1 parent f96f1e5
commit fd818d2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tools/inference/inference_pipeline.ipynb b/tools/inference/inference_pipeline.ipynb
@@ -47,7 +47,7 @@
       "outputs": [],
       "source": [
         "# Required only for colab environments + GPU\n",
-        "!pip install \"jax[cuda12_pip]==0.3.25\" jaxlib==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
+        "!pip install jax==0.3.25 jaxlib==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
         "\n",
         "# Install required libraries\n",
         "!pip install -q dalle-mini orbax==0.0.23\n",