You're ready to start training your model now! Load ViT with [AutoModelForImageClassification]. Specify the number of labels along with the number of expected labels, and the label mappings:

from transformers import AutoModelForImageClassification, TrainingArguments, Trainer
model = AutoModelForImageClassification.from_pretrained(
     checkpoint,
     num_labels=len(labels),
     id2label=id2label,
     label2id=label2id,
 )

At this point, only three steps remain: