def apply_ffn(
        self,
        attention_mask: Optional[torch.ByteTensor],
        m: torch.Tensor,
    ) -> torch.Tensor:
        """Apply feed forward layers to the input.

        Args:
            attention_mask (Optional[torch.ByteTensor]): The attention mask.
            m (torch.Tensor): The input.

        Returns:
            n (torch.Tensor): The output.
        """
        batch_size, seq_len = m.size()[:2]
        indices = None
        if not self.use_pad_tok_in_ffn and attention_mask is not None:
            assert unpad_input is not None
            attention_mask = self.slice_attention_mask(attention_mask, seq_len)
            m, indices, *_ = unpad_input(m, attention_mask)
        n = self.ffn(m)
        if not self.use_pad_tok_in_ffn and attention_mask is not None:
            assert pad_input is not None
            n = pad_input(n, indices, batch_size, seq_len)
        return n
Apply feed forward layers to the input.

        Args:
            attention_mask (Optional[torch.ByteTensor]): The attention mask.
            m (torch.Tensor): The input.

        Returns:
            n (torch.Tensor): The output.
        
llmfoundry/models/layers/blocks.py