def forward(
        self,
        x,
        *,
        mask=None
    ):
        """
        Forward pass through the Mossformer Block.

        Args:
            x (Tensor): Input tensor of shape (batch_size, seq_len, dim).
            mask (Tensor, optional): Attention mask to apply. Defaults to None.

        Returns:
            Tensor: Output tensor after passing through all layers, of shape (batch_size, seq_len, dim).
        """
        ii = 0
        # Iterate through all FLASH attention layers and Gated FSMN blocks
        for flash in self.layers:
            x = flash(x, mask=mask)  # Apply FLASH attention layer
            x = self.fsmn[ii](x)     # Apply corresponding Gated FSMN block
            ii += 1  # Increment index for the Gated FSMN block

        return x  # Return the final output after all layers

        Forward pass through the Mossformer Block.

        Args:
            x (Tensor): Input tensor of shape (batch_size, seq_len, dim).
            mask (Tensor, optional): Attention mask to apply. Defaults to None.

        Returns:
            Tensor: Output tensor after passing through all layers, of shape (batch_size, seq_len, dim).
        
clearvoice/clearvoice/models/mossformer2_ss/mossformer2_block.py