def forward(self, hidden_states):
    hidden_states = hidden_states[:, (self.cls_in_input + self.rel_in_input):, :]
    return self.re_fc(hidden_states)
def forward(self, hidden_states):
    hidden_states = hidden_states[:, (self.cls_in_input + self.rel_in_input):, :]
    return self.re_fc(hidden_states)