def forward(self, hidden_states, attention_mask, segments, segment_index, head_mask=None, output_attentions=False):
    '\n        Input:\n            hidden_states[batch, seg_len, hidden_dim]:\n            attention_mask[batch, seg_len](segment-based)\n            segments[batch, seg_len]:\n            segment_index[batch, seq_len]\n\n        '
    local_states = self.local_attention(hidden_states, attention_mask, head_mask)
    if output_attentions:
        (local_states, attention_probs) = local_states
    input_prefix = (hidden_states.size(1) - segment_index.size(1))
    (seg_states, seg_atten_mask) = self.get_seg_states(local_states[:, input_prefix:, :], segment_index)
    if (self.segment_atten_type == 'multi'):
        seg_states = self.segment_attention(seg_states, seg_atten_mask.unsqueeze(1).unsqueeze(2), head_mask)
        if output_attentions:
            (seg_states, attention_probs) = seg_states
    local_inter_states = self.local_inter_layer(local_states)
    seg_inter_states = self.seg_inter_layer(seg_states)
    aug_seg_inter_states = torch.gather(seg_inter_states, 1, (segments[:, input_prefix:] - 2).view(segments.size(0), (- 1), 1).repeat(1, 1, seg_inter_states.size(2)))
    inter_states = torch.zeros(local_inter_states.size(0), local_inter_states.size(1), self.inter_dim).to(local_inter_states.device)
    inter_states[:, :, :(self.inter_dim // 2)] = local_inter_states
    inter_states[:, input_prefix:, (self.inter_dim // 2):] = aug_seg_inter_states
    inter_states[:, :input_prefix, (self.inter_dim // 2):] = seg_inter_states.sum(dim=1, keepdim=True)
    output_states = self.output(inter_states, hidden_states)
    if output_attentions:
        return (output_states, attention_probs)
    return output_states
Input:
    hidden_states[batch, seg_len, hidden_dim]:
    attention_mask[batch, seg_len](segment-based)
    segments[batch, seg_len]:
    segment_index[batch, seq_len]
def forward(self, hidden_states, attention_mask, segments, segment_index, head_mask=None, output_attentions=False):
    '\n        Input:\n            hidden_states[batch, seg_len, hidden_dim]:\n            attention_mask[batch, seg_len](segment-based)\n            segments[batch, seg_len]:\n            segment_index[batch, seq_len]\n\n        '
    local_states = self.local_attention(hidden_states, attention_mask, head_mask)
    if output_attentions:
        (local_states, attention_probs) = local_states
    input_prefix = (hidden_states.size(1) - segment_index.size(1))
    (seg_states, seg_atten_mask) = self.get_seg_states(local_states[:, input_prefix:, :], segment_index)
    if (self.segment_atten_type == 'multi'):
        seg_states = self.segment_attention(seg_states, seg_atten_mask.unsqueeze(1).unsqueeze(2), head_mask)
        if output_attentions:
            (seg_states, attention_probs) = seg_states
    local_inter_states = self.local_inter_layer(local_states)
    seg_inter_states = self.seg_inter_layer(seg_states)
    aug_seg_inter_states = torch.gather(seg_inter_states, 1, (segments[:, input_prefix:] - 2).view(segments.size(0), (- 1), 1).repeat(1, 1, seg_inter_states.size(2)))
    inter_states = torch.zeros(local_inter_states.size(0), local_inter_states.size(1), self.inter_dim).to(local_inter_states.device)
    inter_states[:, :, :(self.inter_dim // 2)] = local_inter_states
    inter_states[:, input_prefix:, (self.inter_dim // 2):] = aug_seg_inter_states
    inter_states[:, :input_prefix, (self.inter_dim // 2):] = seg_inter_states.sum(dim=1, keepdim=True)
    output_states = self.output(inter_states, hidden_states)
    if output_attentions:
        return (output_states, attention_probs)
    return output_states