def METHOD_NAME(state, batch):
    labels = batch.pop("labels")
    logits = state.apply_fn(**batch, params=state.params, train=False)[0]
    metrics = wrapper.eval_metrics(batch, logits, labels)
    if do_distributed:
        metrics = jax.lax.pmean(metrics, axis_name="batch")
    return logits, metrics