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