examples/research_projects/pplm/pplm_classification_head.py/ClassificationHead/forward
class ClassificationHead:    def forward(self, hidden_state):
        # hidden_state = nn.functional.relu(self.mlp1(hidden_state))
        # hidden_state = self.mlp2(hidden_state)
        logits = self.mlp(hidden_state)
        return logits