    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor,
        head_mask: Optional[torch.Tensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
    ) -> TransformerLayerOutput:
        """
        # Parameters

        hidden_states : `torch.Tensor`
            Shape `batch_size x seq_len x hidden_dim`
        attention_mask : `torch.BoolTensor`, optional
            Shape `batch_size x seq_len`
        head_mask : `torch.BoolTensor`, optional
        encoder_hidden_states : `torch.Tensor`, optional
        encoder_attention_mask : `torch.Tensor`, optional
        output_attentions : `bool`
            Whether to also return the attention probabilities, default = `False`
        """
        attention_outputs = self.attention(
            hidden_states,
            attention_mask,
            head_mask,
            output_attentions=output_attentions,
        )
        attention_output = attention_outputs.hidden_states
        self_attention_probs = attention_outputs.attention_probs
        cross_attention_probs = None

        if encoder_hidden_states is not None:
            assert hasattr(
                self, "cross_attention"
            ), f"If `encoder_hidden_states` are passed, {self} has to be instantiated "
            "with cross-attention layers by setting `config.add_cross_attention=True`"

            cross_attention_outputs = self.cross_attention(
                attention_output,
                attention_mask,
                head_mask,
                encoder_hidden_states,
                encoder_attention_mask,
                output_attentions,
            )
            attention_output = cross_attention_outputs.hidden_states
            cross_attention_probs = cross_attention_outputs.attention_probs

        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output)

        outputs = TransformerLayerOutput(layer_output, self_attention_probs, cross_attention_probs)
        return outputs
# Parameters

hidden_states : `torch.Tensor`
Shape `batch_size x seq_len x hidden_dim`
attention_mask : `torch.BoolTensor`, optional
Shape `batch_size x seq_len`
head_mask : `torch.BoolTensor`, optional
encoder_hidden_states : `torch.Tensor`, optional
encoder_attention_mask : `torch.Tensor`, optional
output_attentions : `bool`
Whether to also return the attention probabilities, default = `False`
allennlp/modules/transformer/transformer_layer.py