code: def train_step(inputs):
      model = tf.keras.Sequential(
          [
              tf.keras.layers.Conv2D(32, 3, activation="relu"),
              tf.keras.layers.MaxPooling2D(),
              tf.keras.layers.Flatten(),
              tf.keras.layers.Dropout(0.1),
              tf.keras.layers.Dense(10, activation="softmax"),
          ]
      )
      loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()
      optimizer = tf.keras.optimizers.SGD()
  
      x, y = inputs
      y_hat = model(x, training=True)
      loss = loss_fn(y, y_hat)
      all_vars = []
      for v in model.trainable_variables:
          all_vars.append(v)
      grads = tf.gradients(loss, all_vars)
      update = optimizer.apply_gradients(zip(grads, all_vars))
  
      return loss, update
  
file_path: autodist/tests/checkpoint/test_keras_saver.py
imports: import os
  import atexit
  import numpy
  import tensorflow
  from multiprocessing import Process
  from autodist import AutoDist
  from autodist.checkpoint.saver import Saver
label_desc: Use image node_id 2 for calling a global function with example usage: train_step(inputs) and returns: loss,  update