def forward(self, hidden_states, attention_mask, head_mask=None, output_attentions=False):
    input_states = hidden_states
    hidden_states = self.attention(hidden_states, attention_mask, head_mask=head_mask)
    if output_attentions:
        (hidden_states, attention_probs) = hidden_states
    output_states = self.output(hidden_states, input_states)
    if output_attentions:
        return (output_states, attention_probs)
    return output_states
def forward(self, hidden_states, attention_mask, head_mask=None, output_attentions=False):
    input_states = hidden_states
    hidden_states = self.attention(hidden_states, attention_mask, head_mask=head_mask)
    if output_attentions:
        (hidden_states, attention_probs) = hidden_states
    output_states = self.output(hidden_states, input_states)
    if output_attentions:
        return (output_states, attention_probs)
    return output_states