Three class prediction but only two output

On the basis of transformers/run_glue.py at master · huggingface/transformers · GitHub I tried to do a run for snli. I thought I included the relevant bits but my prediction matrix only has two columns but the task allows three possible labels.

I used
model = AutoModelForSequenceClassification.from_pretrained(‘bert-base-cased’)
tokenizer = AutoTokenizer.from_pretrained(‘bert-base-cased’)

Copied the mapping from run_glue.py
def preprocess_function(examples):
args = (examples[“premise”], examples[“hypothesis”])
result = tokenizer(*args, truncation=True)
return result
snli = snli.map(
preprocess_function,
batched=True
)

and started the training
training_args = TrainingArguments(…)

trainer = Trainer(
model=model,
args=training_args,
train_dataset=snli[“train”],
compute_metrics=compute_metrics,
tokenizer=tokenizer,
)

But when I predicted trainer.predict(snli[‘test’]) My matrix had a shape of (10000,2) so the whole training was wrong I guess.

What did I miss? Thank you