code: def train_fn(self, xy):
      x, y = xy
      # with tf.GradientTape() as tape:
      loss = self.forward(x, y)
      trainables = [self.emb, self.w1, self.b1, self.w2, self.b2]
      gradients = tf.gradients(loss, trainables)
      # gradients = tape.gradient(loss, trainables)
      # strategy requires users to provide the train_op handle
      train_op = self.optimizer.apply_gradients(
          zip(gradients, trainables)
      )
      return loss, train_op, gradients
  
file_path: autodist/tests/checkpoint/test_partitionedPS_saver.py
imports: import atexit
  import os
  import time
  import tensorflow
  import numpy
  from multiprocessing import Process
  from tensorflow.core.protobuf import config_pb2
  from autodist import AutoDist
  from autodist.checkpoint.saver import Saver
  from autodist.strategy import PartitionedPS
label_desc: Use image node_id 3 for calling the SimpleModel obj's underlying member method code with example usage: obj.train_fn(xy) and returns: loss,  train_op,  gradients