evaluate_model.py

def evaluate_model(state, batch):
    """Evaluate on the validation set."""
    test_imgs, test_lbls = batch
    metrics = eval_step(state, test_imgs, test_lbls)
    metrics = jax.device_get(metrics)
    metrics = jax.tree_map(lambda x: x.item(), metrics)
    return metrics