from transformers import create_optimizer
batch_size = 16
num_epochs = 5
num_train_steps = len(food["train"]) * num_epochs
learning_rate = 3e-5
weight_decay_rate = 0.01
optimizer, lr_schedule = create_optimizer(
     init_lr=learning_rate,
     num_train_steps=num_train_steps,
     weight_decay_rate=weight_decay_rate,
     num_warmup_steps=0,
 )

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