How to set useful id2label and label2id in config.json using Trainer

TL;DR - when training, a config.json is created and I would like to know how to define useful label ids, instead of the default “LABEL_0”.

I’ve defined a dataset loader, which includes a label feature. I’m calling functions such as:

model = ResNetForImageClassification.from_pretrained(...)
train_dataset = load_dataset('custom/dataset/path', split='test')
trainer = CustomTrainer(model=model, ...)

When it saves out checkpoints, it creates a config.json with default id2label labels, such as “0”: “LABEL_0”. How can I pass the real labels so that it ends up in the id2label and label2id in config.json?