From 2facedab298f6bcc50897f4c0cf76354c70a9108 Mon Sep 17 00:00:00 2001
From: Wei Wei <5577741+windmaple@users.noreply.github.com>
Date: Sun, 6 Oct 2024 08:51:53 +0800
Subject: [PATCH] Update distributed_tuning.ipynb

Keras updated ModelParallel API:

https://github.com/keras-team/keras/issues/20039
---
 site/en/gemma/docs/distributed_tuning.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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"