docs: def input_fn(is_training, data_dir, batch_size, num_epochs=1, num_gpus=None,
               dtype=tf.float32):
    """Input function which provides batches for train or eval.
  
    Args:
      is_training: A boolean denoting whether the input is for training.
      data_dir: The directory containing the input data.
      batch_size: The number of samples per batch.
      num_epochs: The number of epochs to repeat the dataset.
      num_gpus: The number of gpus used for training.
      dtype: Data type to use for images/features
  
    Returns:
      A dataset that can be used for iteration.
    """
    mlperf_log.resnet_print(key=mlperf_log.INPUT_ORDER)
    filenames = get_filenames(is_training, data_dir)
    dataset = tf.data.Dataset.from_tensor_slices(filenames)
  
    if is_training:
      # Shuffle the input files
      dataset = dataset.shuffle(buffer_size=_NUM_TRAIN_FILES)
  
    # Convert to individual records
    dataset = dataset.flat_map(tf.data.TFRecordDataset)
  
    return resnet_run_loop.process_record_dataset(
        dataset=dataset,
        is_training=is_training,
        batch_size=batch_size,
        shuffle_buffer=_SHUFFLE_BUFFER,
        parse_record_fn=parse_record,
        num_epochs=num_epochs,
        num_gpus=num_gpus,
        examples_per_epoch=_NUM_IMAGES['train'] if is_training else None,
        dtype=dtype
    )