def forward(self, hidden_states, attention_mask, head_mask=None, output_attentions=False):
    hidden_states = self.attention(hidden_states, attention_mask, head_mask)
    if output_attentions:
        (hidden_states, attention_probs) = hidden_states
    inter_states = self.inter_layer(hidden_states)
    output_states = self.output(inter_states, hidden_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):
    hidden_states = self.attention(hidden_states, attention_mask, head_mask)
    if output_attentions:
        (hidden_states, attention_probs) = hidden_states
    inter_states = self.inter_layer(hidden_states)
    output_states = self.output(inter_states, hidden_states)
    if output_attentions:
        return (output_states, attention_probs)
    return output_states