# Define an Adam optimizer with learning rate 0.001
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001) 
# Define the loss function
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)  

for inputs, labels in dataset:
  with tf.GradientTape() as tape:
    logits = model(inputs)
    # Compute the loss
    loss = loss_fn(labels, logits)  
  # Compute the gradients
  grads = tape.gradient(loss, model.trainable_variables)  
  # Update the model parameters
  optimizer.apply_gradients(zip(grads, model.trainable_variables))

        __
        __