Incorrect use of extend_fsdp_sharding_meta() in cross_fused_attn() (NVIDIA#482)

fixed incorrect of extend_fsdp_sharding_meta() in cross_fused_attn()

Signed-off-by: Alp Dener <[email protected]>
Signed-off-by: Charlene Yang <[email protected]>

Author: denera

transformer_engine/jax/fused_attn.py

@@ -206,7 +206,7 @@ def cross_fused_attn(q: jnp.ndarray,
             tp_dims=([2, 3, None, None], [2]),
             dp_axis_name=dp_axis_name,
             tp_axis_name=tp_axis_name)
-        sharding_meta = extend_fsdp_sharding_meta(sharding_meta, {0: 0, 2: 0})
+        sharding_meta, _ = extend_fsdp_sharding_meta(sharding_meta, {0: 0, 2: 0})
 
         inputs_ = tuple(
             jnp.reshape(x, new_shape) if x is not None else None