code: def main():
      mnist = tf.keras.datasets.mnist
  
      # download and load data
      (x_train, y_train), (x_test, y_test) = mnist.load_data()
      x_train, x_test = x_train / 255.0, x_test / 255.0
  
      # Add a channels dimension
      x_train = x_train[..., tf.newaxis]
      x_test = x_test[..., tf.newaxis]
  
      # create data generator
      train_ds = (
          tf.data.Dataset.from_tensor_slices((x_train, y_train))
          .shuffle(10000)
          .batch(32)
      )
      test_ds = tf.data.Dataset.from_tensor_slices(
          (x_test, y_test)
      ).batch(32)
  
      # create model
      model = MyModel()
  
      # define loss
      loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
      # define optimizer
      optimizer = tf.keras.optimizers.Adam()
  
      # define train_loss and train_accuracy
      train_loss = tf.keras.metrics.Mean(name="train_loss")
      train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
          name="train_accuracy"
      )
  
      # define train_loss and train_accuracy
      test_loss = tf.keras.metrics.Mean(name="test_loss")
      test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
          name="test_accuracy"
      )
  
      # define train function including calculating loss, applying gradient and calculating accuracy
      @tf.function
      def train_step(images, labels):
          with tf.GradientTape() as tape:
              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)
  
      # define test function including calculating loss and calculating accuracy
      @tf.function
      def test_step(images, labels):
          predictions = model(images)
          t_loss = loss_object(labels, predictions)
  
          test_loss(t_loss)
          test_accuracy(labels, predictions)
  
      EPOCHS = 5
  
      for epoch in range(EPOCHS):
          train_loss.reset_states()  # clear history info
          train_accuracy.reset_states()  # clear history info
          test_loss.reset_states()  # clear history info
          test_accuracy.reset_states()  # clear history info
  
          for images, labels in train_ds:
              train_step(images, labels)
  
          for test_images, test_labels in test_ds:
              test_step(test_images, test_labels)
  
          template = "Epoch {}, Loss: {}, Accuracy: {}, Test Loss: {}, Test Accuracy: {}"
          print(
              template.format(
                  epoch + 1,
                  train_loss.result(),
                  train_accuracy.result() * 100,
                  test_loss.result(),
                  test_accuracy.result() * 100,
              )
          )
  
label_desc: Use image node_id 1 for calling a global function with example usage: main() without return types