accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return accuracy, loss_value

  @tf.function
  def distributed_train_step(dist_inputs):
    per_replica_accuracy, per_replica_losses = strategy.run(
        train_step, args=(dist_inputs,))
    accuracy = strategy.reduce(
        tf.distribute.ReduceOp.MEAN, per_replica_accuracy, axis=None)