def METHOD_NAME(iterator):
    def replica_fn(data):
        features, labels = data
        with tf.GradientTape() as tape:
            output = model(tf.expand_dims(features, axis=1))
            loss = keras.losses.MeanSquaredError(
                reduction=losses_utils.ReductionV2.NONE
            )(labels, output)
        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(
            zip(grads, model.trainable_variables)
        )
    strategy.run(replica_fn, args=(next(iterator),))