code: def train_step(images, labels):
      with tf.GradientTape() as tape:
          # training=True is only needed if there are layers with different
          # behavior during training versus inference (e.g. Dropout).
          predictions = model(images)
          loss = loss_object(labels, predictions)
      gradients = tape.gradient(loss, model.trainable_variables)
      optimizer.apply_gradients(
          zip(gradients, model.trainable_variables)
      )
      train_loss(loss)
      train_accuracy(labels, predictions)
  
file_path: BentoML/examples/tensorflow2_native/train.py
imports: import tensorflow
  from tensorflow.keras.layers import Conv2D
  from tensorflow.keras.layers import Dense
  from tensorflow.keras.layers import Flatten
  import bentoml
label_desc: Use image node_id 1 for calling a global function with example usage: train_step(images, labels) without return types