def METHOD_NAME(
    self,
    x: torch.Tensor,
    self_attn_mask: Optional[torch.Tensor] = None,
    self_attn_padding_mask: Optional[torch.Tensor] = None,
):
    """
    LayerNorm is applied either before or after the self-attention/ffn
    modules similar to the original Transformer implementation.
    """
    residual = x
    x, attn = self.self_attn(
        query=x,
        key=x,
        value=x,
        key_padding_mask=self_attn_padding_mask,
        need_weights=False,
        attn_mask=self_attn_mask,
    )
    x = self.dropout_module(x)
    x = residual + x
    x = self.self_attn_layer_norm(x)
    residual = x
    x = self.activation_fn(self.fc1(x))
    x = self.activation_dropout_module(x)
    x = self.fc2(x)
    x = self.dropout_module(x)
    x = residual + x
    x = self.final_layer_norm(x)
    return x, attn