training_loop

def apply_gradient(optimizer, model, x, y):
    """
    Function for computing gradient and updating the weights
    """
    with tf.GradientTape() as tape:
        # Get model prediction and compute the loss
        logits = model(x)
        loss_value = loss_object(y_true=y, y_pred=logits)
    # Calculate the gradient using tape.gradient and then update the model weights using our optimizer
    gradients = tape.gradient(loss_value, model.trainable_weights)
    optimizer.apply_gradients(zip(gradients, model.trainable_weights))  
    return logits, loss_value