update cuda version of jax dependency in mt3 colab

PiperOrigin-RevId: 636999626
magenta · May 25, 2024 · fdc09df · fdc09df
1 parent ef4f9b2
commit fdc09df
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[cuda11_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_local] 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",