eval_step.py

@jax.jit
def eval_step(state, batch):
    images, labels = batch
    logits = CNN().apply({'params': state.params}, images)
    return compute_metrics(logits=logits, labels=labels)