diff --git a/notebooks/unit5/unit5.ipynb b/notebooks/unit5/unit5.ipynb
index a19498b..aeed808 100644
--- a/notebooks/unit5/unit5.ipynb
+++ b/notebooks/unit5/unit5.ipynb
@@ -200,7 +200,7 @@
       "source": [
         "%%capture\n",
         "# Clone the repository\n",
-        "!git clone --depth 1 --branch hf-integration https://github.com/huggingface/ml-agents"
+        "!git clone --depth 1 --branch hf-integration-save https://github.com/huggingface/ml-agents"
       ]
     },
     {
@@ -865,4 +865,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 0
-}
\ No newline at end of file
+}