code: def create_input_fn(prefix, batch_size, augmentation=None):
      
      # Check if the version is too old for dataset api to work better than manually loading data.
      if tf.__version__.startswith("1.6") or tf.__version__.startswith("1.5") or tf.__version__.startswith("1.4") \
              or tf.__version__.startswith("1.3") or tf.__version__.startswith("1.2") \
              or tf.__version__.startswith("1.1") or tf.__version__.startswith("1.0"):
          def input_fn():
              with tf.variable_scope("input_pipeline"):
                  return _read_data_legacy(prefix, batch_size)
          return input_fn
      else:
          def input_fn():
              with tf.variable_scope("input_pipeline"):
                  return _read_data(prefix, batch_size, augmentation)
          return input_fn
docs: Loads a dataset.
  
      :param prefix: The path prefix as defined in the write data method.
      :param batch_size: The batch size you want for the tensors.
      :param augmentation: An augmentation function.
      :return: An input function for a tf estimator.