LayoutLMV3 mapping numerical predictions to class labels

Hello,

I trained a LayoutLMV3 model to predict document class. I used the hugging face datasets library to prepare and load a custom dataset. I created a dictionary id2label to transform numerical to categorical labels in my model training. However, I still get the predictions as numerical labels. Is there any solution to this issue except running the id2label again for predictions?

Thank you in advance