final fix

Author: parambole

MaxText/convert_qwen3_moe.py

@@ -272,4 +272,4 @@ def main(args):
   parser.add_argument("--use-zarr3", type=str2bool, default=True, help="Use Zarr3 format for saving.")
 
   parsed_args = parser.parse_args()
-  main(parsed_args)
+  main(parsed_args)

MaxText/layers/moe.py

@@ -382,7 +382,6 @@ def expert_group_mask(self, gate_logits: jax.Array) -> jax.Array:
 
     # Mask selected groups so that only those experts are considered.
     group_mask = jax.nn.one_hot(group_idx, num_classes=self.config.n_routing_groups, dtype=jnp.float32)
-    group_mask = jax.nn.one_hot(group_idx, num_classes=self.config.n_routing_groups, dtype=jnp.float32)
     group_mask = jnp.sum(group_mask, axis=-2)
 
     # Apply masks and get top-k indices.