for _ in range(4):
    with tf.GradientTape() as tape:
        predictions = model(FEATURES, training=True)
        loss = tf.reduce_mean(
            tf.keras.losses.binary_crossentropy(LABELS, predictions)
        )
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    training_step.assign_add(1)