diff --git a/src/maxtext/examples/sft_llama3_demo_tpu.ipynb b/src/maxtext/examples/sft_llama3_demo_tpu.ipynb
index d4fba92eed..9bb02d5bc9 100644
--- a/src/maxtext/examples/sft_llama3_demo_tpu.ipynb
+++ b/src/maxtext/examples/sft_llama3_demo_tpu.ipynb
@@ -155,7 +155,11 @@
     "    notebook_login()\n",
     "else:\n",
     "    from huggingface_hub import login\n",
-    "    login()"
+    "    hf_token = os.environ.get(\"HF_TOKEN\", \"\")\n",
+    "    if hf_token:\n",
+    "        login(token=hf_token)\n",
+    "    else:\n",
+    "        login()"
    ]
   },
   {