  def train_step(optimizer, batch):
    """Train for a single step."""
    def loss_fn(model):
-@@ -121,18 +139,23 @@ def train_step(optimizer, batch):
+@@ -120,19 +136,25 @@ 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)
++  grad = lax.pmean(grad, 'batch')
    optimizer = optimizer.apply_gradient(grad)
 -  metrics = compute_metrics(logits, batch['label'])
 +  metrics = compute_metrics(logits, batch['label'], sharded=True)

       'accuracy': accuracy,
   }
+  if sharded:
+    metrics = pmean(metrics)
   return metrics
 
 

   train_ds_size = len(train_ds['image'])
   steps_per_epoch = train_ds_size // batch_size
 
@@ -141,7 +163,7 @@ def train_epoch(optimizer, train_ds, batch_size, epoch, rng):
   perms = perms.reshape((steps_per_epoch, batch_size))
   batch_metrics = []
   for perm in perms:
