'Target 3 is out of bounds' error with Huggingface Trainer.train()

I have 3 labels, so the num_labels is set to 3. I have also double-checked my data and ensured there are no blank values. However, when I train my model using trainer.train(), after a while, it throws this error: ‘Target 3 is out of bounds.’ After searching online, I found that it could be an index issue. However, I don’t know how to change the index of my labels to [0, 1, 2], so could someone please explain the code and what could be causing the error?

Here is the code snippet below with the num_labels as a parameter for the pre-trained model:

from transformers import ViTFeatureExtractor
model_name_or_path = ‘google/vit-base-patch16-224-in21k’
feature_extractor = ViTFeatureExtractor.from_pretrained(model_name_or_path)
from transformers import ViTForImageClassification
labels = dataset[‘train’].features[‘labels’].names
model = ViTForImageClassification.from_pretrained(
model_name_or_path,
num_labels=len(labels)
)

model.train()