code: def main(_):
      data_path = os.path.join(
          FLAGS.datadir, "training-monolingual.tokenized.shuffled/*"
      )
      distribute_batch_size = (
          FLAGS.batch_size * autodist._resource_spec.num_gpus
      )
      with tf.Graph().as_default(), autodist.scope():
          train_dataset = gen_lm1b_train_dataset(
              data_path, FLAGS.num_steps
          )
          train_dataset = train_dataset.batch(FLAGS.batch_size)
          train_iterator = (
              train_dataset.make_one_shot_iterator().get_next()
          )
  
          model = language_model.LM(FLAGS.num_steps)
          # TODO (Hao): need to improve this.
          train_step = autodist.function(model.train_step)
  
          prev_time = time.time()
          for local_step in range(FLAGS.max_steps):
              loss, _ = train_step(train_iterator)
              if local_step % FLAGS.log_frequency == 0:
                  cur_time = time.time()
                  elapsed_time = cur_time - prev_time
                  num_words = (
                      distribute_batch_size * FLAGS.log_frequency
                  )
                  wps = float(num_words) / elapsed_time
                  print(
                      "Iteration %d, time = %.2fs, wps = %.0f, train loss = %.4f"
                      % (local_step, cur_time - prev_time, wps, loss)
                  )
                  prev_time = cur_time
  
file_path: autodist/examples/lm1b/lm1b_train.py
imports: import glob
  import os
  import time
  import tensorflow
  from absl import app
  from absl import flags
  import language_model
  from autodist import AutoDist
  from autodist.strategy import PS, PSLoadBalancing, PartitionedPS, AllReduce, Parallax
label_desc: Use image node_id 2 for calling a global function with example usage: main(_) without return types