Num_labels creates an error for some models


I am training a classifier for three classes (bad, medium, good) using distilbert-base-uncased-finetuned-sst-2-english but after saving the model to disk and re-loading it I am getting a strange error

from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
model = TFAutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")

model.from_pretrained('finetuned', num_labels = 3)

ValueError: cannot reshape array of size 1536 into shape (768,3)

This approach worked with bert-base-uncased. Am I doing something wrong here?

With the latest version installed, you need to add ignore_mismatched_sizes=True to your from_pretrained call for this to work.
Otherwise, you try to load a model with 2 labels inside a model with 3 labels, and you get mismatched sizes like that.


thanks @sgugger, life saving tip! :pray:

1 Like