| 93 | |
| 94 | |
| 95 | class ClsMixin(BaseMixin): |
| 96 | def __init__(self, hidden_size, num_classes): |
| 97 | super().__init__() |
| 98 | self.classifier = nn.Linear(hidden_size, num_classes) |
| 99 | |
| 100 | def final_forward(self, logits, **kw_args): |
| 101 | logits = self.classifier(logits[:, 0]) |
| 102 | return logits |
| 103 | |
| 104 | |
| 105 | class ViTModel(BaseModel): |