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

     optimizer, metrics = train_step(optimizer, batch)
     batch_metrics.append(metrics)
 
@@ -191,7 +206,8 @@ def train(train_ds, test_ds):
     rng, input_rng = random.split(rng)
     optimizer, train_metrics = train_epoch(
         optimizer, train_ds, batch_size, epoch, input_rng)
-    loss, accuracy = eval_model(optimizer.target, test_ds)
+    model = jax_utils.unreplicate(optimizer.target)  # Fetch from 1st device
+    loss, accuracy = eval_model(model, test_ds)
     logging.info('eval epoch: %d, loss: %.4f, accuracy: %.2f',
                  epoch, loss, accuracy * 100)
     summary_writer.scalar('train_loss', train_metrics['loss'], epoch)
diff --git a/examples/mnist/train_test.py b/examples/mnist/train_test.py
index 429bda4..3ebf534 100644
--- a/examples/mnist/train_test.py
+++ b/examples/mnist/train_test.py
@@ -21,6 +21,8 @@ import train
 import jax
 from jax import random
 
+from flax import jax_utils
+
 # Parse absl flags test_srcdir and test_tmpdir.
 jax.config.parse_flags_with_absl()
 
@@ -36,12 +38,13 @@ class TrainTest(absltest.TestCase):
     # test single train step.
     optimizer, train_metrics = train.train_step(
         optimizer=optimizer,
-        batch={k: v[:batch_size] for k, v in train_ds.items()})
+        batch=train.shard({k: v[:batch_size] for k, v in train_ds.items()}))
     self.assertLessEqual(train_metrics['loss'], 2.302)
     self.assertGreaterEqual(train_metrics['accuracy'], 0.0625)
 
     # Run eval model.
-    loss, accuracy = train.eval_model(optimizer.target, test_ds)
+    model = jax_utils.unreplicate(optimizer.target) # Fetch from 1st device
+    loss, accuracy = train.eval_model(model, test_ds)
     self.assertLess(loss, 2.252)
     self.assertGreater(accuracy, 0.2597)
 

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