fix norm sharding

google · May 2, 2024 · bc36642 · bc36642
1 parent fcf48fe
commit bc36642
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)