def METHOD_NAME(input_context):
    image_size = (608, 608)
    device_id = input_context.input_pipeline_id
    num_threads = input_context.num_input_pipelines
    if pipeline == 'dali-gpu' or pipeline == 'dali-cpu':
        with tf.device("/gpu:{}".format(input_context.input_pipeline_id)):
            yolo = YOLOv4Pipeline(
                file_root, annotations,
                batch_size, image_size, num_threads, device_id, seed,
                use_gpu=pipeline == 'dali-gpu',
                is_training=is_training,
                use_mosaic=use_mosaic
            )
            return yolo.dataset()
    if pipeline == 'numpy':
        yolo = YOLOv4PipelineNumpy(
            file_root, annotations,
            batch_size, image_size, num_threads, device_id, seed,
            is_training=is_training,
            use_mosaic=use_mosaic
        )
        return yolo.dataset()