def METHOD_NAME(iterator):
    """The step function for one training step."""
    def step_fn(inputs):
        """The computation to run on each replica(GPU)."""
        features, labels = inputs
        with tf.GradientTape() as tape:
            pred = model(features, training=True)
            loss = keras.losses.binary_crossentropy(labels, pred)
            loss = tf.nn.compute_average_loss(loss)
        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(
            list(zip(grads, model.trainable_variables))
        )
        actual_pred = tf.cast(tf.greater(pred, 0.5), tf.int64)
        accuracy.update_state(labels, actual_pred)
    distribution.run(step_fn, args=(next(iterator),))