  def create_optimizer(model, learning_rate, beta):
    optimizer_def = optim.Momentum(learning_rate=learning_rate, beta=beta)
    optimizer = optimizer_def.create(model)
-@@ -109,7 +114,7 @@ def compute_metrics(logits, labels):
+@@ -82,7 +87,7 @@ def compute_metrics(logits, labels):
    return metrics
  
  

   return summary['loss'], summary['accuracy']
 
 
@@ -183,16 +195,19 @@ def train(train_ds, test_ds):
 
   summary_writer = tensorboard.SummaryWriter(model_dir)
 
-  rng, init_rng = random.split(rng)
+  rng, init_rng = random.split(rng, jax.device_count())
   model = create_model(init_rng)
   optimizer = create_optimizer(model, FLAGS.learning_rate, FLAGS.momentum)
 
+  test_ds = jax_utils.replicate(test_ds)
+

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