def forward(self, hidden_states):
    pooled = hidden_states[:, self.pool_dim, :]
    for rel_fc in self.rel_fcs:
        pooled = rel_fc(pooled)
    return pooled
def forward(self, hidden_states):
    pooled = hidden_states[:, self.pool_dim, :]
    for rel_fc in self.rel_fcs:
        pooled = rel_fc(pooled)
    return pooled