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"