diff --git a/mt3/colab/music_transcription_with_transformers.ipynb b/mt3/colab/music_transcription_with_transformers.ipynb
index 169ea64..21b599c 100644
--- a/mt3/colab/music_transcription_with_transformers.ipynb
+++ b/mt3/colab/music_transcription_with_transformers.ipynb
@@ -87,7 +87,7 @@
         "\n",
         "# downgrade jax\n",
         "# TODO(iansimon): remove this once we can use the latest t5x\n",
-        "!pip install jax==0.3.14\n",
+        "!pip install jax==0.3.15\n",
         "\n",
         "# copy checkpoints\n",
         "!gsutil -q -m cp -r gs://mt3/checkpoints .\n",