Then, load ViT with [TFAutoModelForImageClassification] along with the label mappings:

from transformers import TFAutoModelForImageClassification
model = TFAutoModelForImageClassification.from_pretrained(
     checkpoint,
     id2label=id2label,
     label2id=label2id,
 )

Convert your datasets to the tf.data.Dataset format using the [~datasets.Dataset.to_tf_dataset] and your data_collator: