def METHOD_NAME(test_dist_dataset):
  """Testing Step."""
  def step_fn(images, labels):
    logits = model(images, training=False)
    loss = compute_loss(logits, labels)
    test_loss.update_state(loss)
    test_accuracy.update_state(labels, logits)
  for images, labels in test_dist_dataset:
    strategy.experimental_run_v2(step_fn, args=(images, labels))