Then, load SegFormer with [TFAutoModelForSemanticSegmentation] along with the label mappings, and compile it with the
optimizer. Note that Transformers models all have a default task-relevant loss function, so you don't need to specify one unless you want to:

from transformers import TFAutoModelForSemanticSegmentation
model = TFAutoModelForSemanticSegmentation.from_pretrained(
     checkpoint,
     id2label=id2label,
     label2id=label2id,
 )
model.compile(optimizer=optimizer)  # No loss argument!