def loss_fn(params):
    variables = {"params": params}
    variables.update(state.model_state)
    logits, new_model_state = model(train=True).apply(
        variables,
        batch["image"],
        mutable=["batch_stats"],
        rngs={"dropout": rng})
    loss = jnp.mean(
        losses.softmax_cross_entropy_loss(logits=logits, labels=batch["label"]))
    if weight_decay > 0:
      weight_penalty_params = jax.tree_leaves(variables["params"])
      weight_l2 = sum(
          [jnp.sum(x**2) for x in weight_penalty_params if x.ndim > 1])
      weight_penalty = weight_decay * 0.5 * weight_l2
      loss = loss + weight_penalty
    new_model_state = dict(new_model_state)
    return loss, (new_model_state, logits)