    train_ds_size = len(train_ds['image'])
    steps_per_epoch = train_ds_size // batch_size
  
-@@ -141,7 +164,7 @@ def train_epoch(optimizer, train_ds, batch_size, epoch, rng):
+@@ -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:

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

 def train_step(optimizer, batch):
   """Train for a single step."""
   def loss_fn(model):
@@ -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)
