Merge pull request #637 from google:lizhiyu/miss_sharding

PiperOrigin-RevId: 630121986
google · May 2, 2024 · 2ac0af9 · 2ac0af9
2 parents fcf48fe + bc36642
commit 2ac0af9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/MaxText/layers/models.py b/MaxText/layers/models.py
@@ -315,7 +315,7 @@ def __call__(
  weight_dtype=cfg.weight_dtype,
  name="decoder_norm",
  epsilon=cfg.normalization_layer_epsilon,
- kernel_axes=("embed",),
+ kernel_axes=("norm",),
  )(y)
  y = nn.Dropout(rate=cfg.dropout_rate, broadcast_dims=(-2,))(y, deterministic=deterministic)