def METHOD_NAME(self, data):
    tf.print(train_step_message)
    x, y = data
    with tf.GradientTape() as tape:
        y_pred = self(x)
        loss = self.compiled_loss(y, y_pred)
    gradients = tape.gradient(loss, self.trainable_variables)
    self.optimizer.apply_gradients(
        zip(gradients, self.trainable_variables)
    )
    return {}