How do I change the classification head of a model?

🤗Transformers
12

Well, I tried that and added two lines to set the correct number of labels:

model.classifier = nn.Linear(786,1)
model.num_labels = 2
model.config.num_labels = 2

printing the model shows that this worked.

  (dropout): Dropout(p=0.1, inplace=False)
  (classifier): Linear(in_features=786, out_features=1, bias=True)

However, running this with the trainer class results in this error:

mat1 dim 1 must match mat2 dim 0

I suspect that I missed something or broke some clever autoconfiguration. Thats why I wrote my question here.

3 Likes