code: def input_fn(dataset, mode, batch_size, num_epoch=None):
      """
      params: dataset, tf.data.Dataset
      params: mode, learning phase
      params: batch size
      params: num of epoch
      """
      if mode == utils.TRAIN:
          _, num_gpus = utils.gpu_device_names()
          per_device_batch_size = utils.per_device_batch_size(
              batch_size, num_gpus
          )
      else:
          # using one device to eval or infer,
          # otherwise will drop reminder samples, e.g. 32 batch with 3 gpus
          per_device_batch_size = batch_size
          num_epoch = 1
  
      logging.info(
          "Learning Phase: {}, Total Batch size:{}, Per device batch size: {}".format(
              mode, batch_size, per_device_batch_size
          )
      )
  
      def _input_fn():
          return dataset(mode, per_device_batch_size, num_epoch)
  
      return _input_fn
  
file_path: delta/delta/data/utils/common_utils.py
imports: import re
  import json
  import numpy
  from absl import logging
  import delta.compat
  import subprocess
  from delta import utils
label_desc: Use image node_id 1 for calling a global function with example usage: input_fn(dataset, mode, batch_size, num_epoch) and returns: _input_fn, dataset