Update distribution.py (#1891)

Fix error after keras-team/keras@0ee33b5
keras-team · Jul 20, 2024 · 63fe4ea · 63fe4ea
1 parent a5c8a58
commit 63fe4ea
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/guides/distribution.py b/guides/distribution.py
@@ -194,7 +194,7 @@
 layout_map["d2/output"] = ("data", None)
 
 model_parallel = keras.distribution.ModelParallel(
-    mesh_2d, layout_map, batch_dim_name="data"
+    layout_map, batch_dim_name="data"
 )
 
 keras.distribution.set_distribution(model_parallel)