diff --git a/site/en/gemma/docs/distributed_tuning.ipynb b/site/en/gemma/docs/distributed_tuning.ipynb
index 6d66e547f..8e40ffe28 100644
--- a/site/en/gemma/docs/distributed_tuning.ipynb
+++ b/site/en/gemma/docs/distributed_tuning.ipynb
@@ -423,7 +423,7 @@
       ],
       "source": [
         "model_parallel = keras.distribution.ModelParallel(\n",
-        "    device_mesh, layout_map, batch_dim_name=\"batch\")\n",
+        "    layout_map=layout_map, batch_dim_name=\"batch\")\n",
         "\n",
         "keras.distribution.set_distribution(model_parallel)\n",
         "gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset(\"gemma_7b_en\")\n"