def METHOD_NAME(defun):
    optimizer = keras.optimizers.legacy.adam.Adam()
    model = test_utils.get_model_from_layers(
        [LayerWithLosses()], input_shape=(10,)
    )
    def train_step(x):
        with tf.GradientTape() as tape:
            model(x)
            assert len(model.losses) == 2
            loss = tf.reduce_sum(model.losses)
        gradients = tape.gradient(loss, model.trainable_weights)
        optimizer.apply_gradients(zip(gradients, model.trainable_weights))
        return loss
    if defun:
        train_step = tf.function(train_step)
    x = tf.ones((10, 10))
    return train_step(x)