examples/research_projects/pplm/pplm_classification_head.py/ClassificationHead/__init__
class ClassificationHead:    def __init__(self, class_size, embed_size):
        super().__init__()
        self.class_size = class_size
        self.embed_size = embed_size
        # self.mlp1 = nn.Linear(embed_size, embed_size)
        # self.mlp2 = (nn.Linear(embed_size, class_size))
        self.mlp = nn.Linear(embed_size, class_size)