code: def train_step(x, y, net, optimizer, loss_summary, accuracy_summary):
      """
      x: input
      y: true label
      net: model object
      optim: optimizer
      loss_summary: summary writer for loss
      acc_summary: summary writer for accuracy
      """
      with tf.GradientTape() as tape:
          z = net(x)
          loss = tf.keras.losses.sparse_categorical_crossentropy(
              y_true=y, y_pred=z, from_logits=True, axis=-1
          )
          loss = tf.reduce_mean(loss)
      grads = tape.gradient(loss, net.trainable_variables)
      optimizer.apply_gradients(zip(grads, net.trainable_variables))
  
      # instrument loss
      loss_summary(loss)
  
      # instrument accuracy
      accuracy_summary(y, z)
      return
  
file_path: sagemaker-python-sdk/tests/data/tensorflow_mnist/mnist_v2.py
imports: import argparse
  import numpy
  import os
  import json
  import tensorflow
  from tensorflow.keras.layers import Conv2D, BatchNormalization, Dense, Flatten
label_desc: Use image node_id 1 for calling a global function with example usage: train_step(x, y, net, optimizer, loss_summary, accuracy_summary) without return types