def forward(self, hidden_states):
    first_token_tensor = hidden_states[:, 0]
    pooled_output = self.fc(first_token_tensor)
    pooled_output = self.activation(pooled_output)
    return pooled_output
def forward(self, hidden_states):
    first_token_tensor = hidden_states[:, 0]
    pooled_output = self.fc(first_token_tensor)
    pooled_output = self.activation(pooled_output)
    return pooled_output