def METHOD_NAME(defun):
    optimizer = keras.optimizers.legacy.adadelta.Adadelta()
    model = test_utils.get_model_from_layers(
        [
            keras.layers.BatchNormalization(momentum=0.9),
            keras.layers.Dense(
                1, kernel_initializer="zeros", activation="softmax"
            ),
        ],
        input_shape=(10,),
    )
    def train_step(x, y):
        with tf.GradientTape() as tape:
            y_pred = model(x, training=True)
            loss = keras.losses.binary_crossentropy(y, y_pred)
        gradients = tape.gradient(loss, model.trainable_weights)
        optimizer.apply_gradients(zip(gradients, model.trainable_weights))
        return loss, model(x, training=False)
    if defun:
        train_step = tf.function(train_step)
    x, y = tf.ones((10, 10)), tf.ones((10, 1))
    return train_step(x, y)