  def train_step(optimizer, batch):
    """Train for a single step."""
    def loss_fn(model):
-@@ -118,8 +127,10 @@ def train_step(optimizer, batch):
+@@ -91,8 +100,10 @@ def train_step(optimizer, batch):
      return loss, logits
    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
    (_, logits), grad = grad_fn(optimizer.target)

   return optimizer
 
 
@@ -99,6 +103,11 @@ def cross_entropy_loss(logits, labels):
   return -jnp.mean(jnp.sum(onehot(labels) * logits, axis=-1))
 
 

   return optimizer, metrics
 
 
@@ -104,6 +115,9 @@ def eval_step(model, batch):
 
 def train_epoch(optimizer, train_ds, batch_size, epoch, rng):
   """Train for a single epoch."""
