As a final step, create a batch of examples using [DefaultDataCollator]:

from transformers import DefaultDataCollator
data_collator = DefaultDataCollator()

Train the model
You’re ready to start training your model now! Load ViLT with [ViltForQuestionAnswering]. Specify the number of labels 
along with the label mappings:

from transformers import ViltForQuestionAnswering
model = ViltForQuestionAnswering.from_pretrained(model_checkpoint, num_labels=len(id2label), id2label=id2label, label2id=label2id)