def moe_fused_init_setup(
    ffn: nn.Module,
):
    """Attach the _stack_dim attribute to the FFN.

    Args:
        ffn (nn.Module): The FFN module.
    """
    ffn.experts.mlp._stack_dim = 0  # type: ignore
Attach the _stack_dim attribute to the FFN.

    Args:
        ffn (nn.Module): The FFN module.
    
llmfoundry/models/layers/ffn.py