  def compute_metrics(logits, labels):
    loss = cross_entropy_loss(logits, labels)
    accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)
-@@ -109,7 +118,7 @@ def compute_metrics(logits, labels):
+@@ -82,7 +91,7 @@ def compute_metrics(logits, labels):
    return metrics
  
  

diff --git a/examples/mnist/train.py b/examples/mnist/train.py
index 9c21985..fed5421 100644
--- a/examples/mnist/train.py
+++ b/examples/mnist/train.py
@@ -23,6 +23,9 @@ from absl import app
 from absl import flags
 from absl import logging
 
+import functools
+
+from flax import jax_utils
 from flax import nn
 from flax import optim
 from flax.metrics import tensorboard
@@ -87,6 +90,7 @@ def create_model(key):
 def create_optimizer(model, learning_rate, beta):
   optimizer_def = optim.Momentum(learning_rate=learning_rate, beta=beta)
   optimizer = optimizer_def.create(model)

   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."""
