diff --git a/mt3/colab/music_transcription_with_transformers.ipynb b/mt3/colab/music_transcription_with_transformers.ipynb
index f885e47..f6fa61c 100644
--- a/mt3/colab/music_transcription_with_transformers.ipynb
+++ b/mt3/colab/music_transcription_with_transformers.ipynb
@@ -90,6 +90,9 @@
         "# copy soundfont (originally from https://sites.google.com/site/soundfonts4u)\n",
         "!gsutil -q -m cp gs://magentadata/soundfonts/SGM-v2.01-Sal-Guit-Bass-V1.3.sf2 .\n",
         "\n",
+        "#Install the correct version of the jax package\n",
+        "!pip install jax==0.3.15\n",
+        "\n",
         "import json\n",
         "import IPython\n",
         "\n",