def METHOD_NAME(images, labels):
    # with tf.GradientTape() as tape:
    predictions = model(images, training=True)
    loss = loss_object(labels, predictions)
    grads = optimizer.get_gradients(loss, model.trainable_variables)
    return grads, loss, predictions