code: def forward(self, x, y):
      # embedding layer
      x = tf.nn.embedding_lookup(self.emb, x)
      # global average pool
      x = tf.math.reduce_mean(x, axis=1)
      # dense
      x = tf.linalg.matmul(x, self.w1) + self.b1
      x = tf.nn.relu(x)
      logits = tf.linalg.matmul(x, self.w2) + self.b2  # logits
      logits = tf.squeeze(logits)
      loss = tf.nn.sigmoid_cross_entropy_with_logits(
          labels=y, logits=logits
      )
      loss = tf.reduce_mean(loss)
      return loss
  
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 2 for calling the SimpleModel obj's underlying member method code with example usage: obj.forward(x, y) and returns: loss