switch jax version "extra" to cuda12

PiperOrigin-RevId: 677009602
magenta · Sep 23, 2024 · e9a31b5 · e9a31b5
1 parent 36359ad
commit e9a31b5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mt3/colab/music_transcription_with_transformers.ipynb b/mt3/colab/music_transcription_with_transformers.ipynb
@@ -67,7 +67,7 @@
         "# install mt3\n",
         "!git clone --branch=main https://github.com/magenta/mt3\n",
         "!mv mt3 mt3_tmp; mv mt3_tmp/* .; rm -r mt3_tmp\n",
-        "!python3 -m pip install jax[cuda12_local] nest-asyncio pyfluidsynth==1.3.0 -e . -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
+        "!python3 -m pip install jax[cuda12] nest-asyncio pyfluidsynth==1.3.0 -e . -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
         "\n",
         "# copy checkpoints\n",
         "!gsutil -q -m cp -r gs://mt3/checkpoints .\n",