import evaluate
metric = evaluate.load("mean_iou")

Then create a function to [~evaluate.EvaluationModule.compute] the metrics. Your predictions need to be converted to
logits first, and then reshaped to match the size of the labels before you can call [~evaluate.EvaluationModule.compute]: