code: def main(_):
      (train_data, train_labels), (
          test_data,
          test_labels,
      ) = tf.keras.datasets.imdb.load_data(num_words=vocab_size)
      train_data = tf.keras.preprocessing.sequence.pad_sequences(
          train_data, value=0, padding="post", maxlen=256
      )
      test_data = tf.keras.preprocessing.sequence.pad_sequences(
          test_data, value=0, padding="post", maxlen=256
      )
      train_labels = train_labels.astype(np.float32)
      with tf.Graph().as_default(), autodist.scope():  # AutoDist code
          my_iterator = (
              tf.compat.v1.data.Dataset.from_tensor_slices(
                  (train_data, train_labels)
              )
              .shuffle(25000)
              .batch(batch_size)
              .repeat()
              .make_one_shot_iterator()
              .get_next()
          )
          # my_iterator = MyIterator().get_next()
          model = SimpleModel()
          prev_time = time.time()
          # fetch train_op and loss
          loss_fn, train_op, gradients = model.train_fn(my_iterator)
          sess = autodist.create_distributed_session()
          for local_step in range(max_steps):
              loss, _ = sess.run(
                  fetches=[loss_fn, train_op],
                  options=config_pb2.RunOptions(
                      trace_level=config_pb2.RunOptions.NO_TRACE
                  ),
              )
              if local_step % log_frequency == 0:
                  cur_time = time.time()
                  elapsed_time = cur_time - prev_time
                  num_sentences = batch_size * log_frequency
                  wps = float(num_sentences) / elapsed_time
                  print(
                      "Iteration %d, time = %.2fs, wps = %.0f, train loss = %.4f"
                      % (local_step, cur_time - prev_time, wps, loss)
                  )
                  prev_time = cur_time
          print(sess.run(model.emb))
          print(sess.run(gradients[0]))
      print("ending...")
  
file_path: autodist/examples/sentiment_classifier.py
imports: import os
  import time
  import tensorflow
  import numpy
  from absl import app
  from tensorflow.core.protobuf import config_pb2
  from autodist import AutoDist
  from autodist.strategy import PS, PSLoadBalancing, PartitionedPS, AllReduce, Parallax
label_desc: Use image node_id 1 for calling a global function with example usage: main(_) without return types