code: def _get_input_iterator(input_fn, strategy):
      """Returns distributed dataset iterator."""
      # When training with TPU pods, datasets needs to be cloned across
      # workers. Since Dataset instance cannot be cloned in eager mode, we instead
      # pass callable that returns a dataset.
      if not callable(input_fn):
          raise ValueError(
              "`input_fn` should be a closure that returns a dataset."
          )
      if not isinstance(strategy, tf.distribute.Strategy):
          dataset = input_fn()
          iterator = tf.compat.v1.data.make_one_shot_iterator(
              input_fn()
          )
      else:
          iterator = iter(
              strategy.experimental_distribute_datasets_from_function(
                  input_fn
              )
          )
      return iterator
  
file_path: autodist/examples/benchmark/utils/bert_utils.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import json
  import os
  import sys
  import tensorflow
  import time
  from absl import logging
label_desc: Use image node_id 1 for calling a global function with example usage: _get_input_iterator(input_fn, strategy) and returns: iterator