                 epoch_metrics_np['loss'], epoch_metrics_np['accuracy'] * 100)
  
    return optimizer, epoch_metrics_np
-@@ -158,7 +170,7 @@ def train_epoch(optimizer, train_ds, batch_size, epoch, rng):
+@@ -131,7 +143,7 @@ def train_epoch(optimizer, train_ds, batch_size, epoch, rng):
  def eval_model(model, test_ds):
    metrics = eval_step(model, test_ds)
    metrics = jax.device_get(metrics)

 def train_step(optimizer, batch):
   """Train for a single step."""
   def loss_fn(model):
@@ -123,7 +128,7 @@ def train_step(optimizer, batch):
   return optimizer, metrics
 
 

diff --git a/examples/mnist/mnist_lib.py b/examples/mnist/mnist_lib.py
index 63af6e8..3901652 100644
--- a/examples/mnist/mnist_lib.py
+++ b/examples/mnist/mnist_lib.py
@@ -20,6 +20,8 @@ The data is loaded using tensorflow_datasets.
 
 from absl import logging
 
+import functools
+
 import jax
 from jax import random
 
@@ -29,6 +31,7 @@ import numpy as onp
 
 import tensorflow_datasets as tfds
 
+from flax import jax_utils
 from flax import nn
 from flax import optim
 from flax.metrics import tensorboard
@@ -51,12 +54,14 @@ class CNN(nn.Module):
     return x
 
 
