def METHOD_NAME(inputs):
  images, labels = inputs
  logits = model(images, training=False)
  loss = compute_loss(logits, labels)
  update_loss = test_loss.update_state(loss)
  update_accuracy = test_accuracy.update_state(labels, logits)
  with tf.control_dependencies([update_loss, update_accuracy]):
    return tf.identity(loss)