code: def get_dataset(
      args, total_batch_size, is_training, params, strategy=None
  ):
      pipeline = args.pipeline_type
  
      if (
          strategy
          and not is_training
          and pipeline == PipelineType.dali_gpu
      ):
          strategy = None
          pipeline = PipelineType.dali_cpu
  
      if pipeline in [PipelineType.tensorflow, PipelineType.synthetic]:
          from pipeline.tf.dataloader import InputReader
  
          if args.input_type != InputType.tfrecord:
              raise ValueError(
                  "tensorflow and syntax pipelines are only compatible with tfrecord input type :<"
              )
  
          if is_training:
              file_pattern = args.train_file_pattern
          else:
              file_pattern = (
                  args.eval_file_pattern or args.train_file_pattern
              )
  
          dataset = InputReader(
              params,
              file_pattern,
              is_training=is_training,
              use_fake_data=(pipeline == PipelineType.synthetic),
          ).get_dataset(total_batch_size)
      elif strategy:
          from pipeline.dali.efficientdet_pipeline import (
              EfficientDetPipeline,
          )
  
          if pipeline == PipelineType.dali_cpu:
              raise ValueError(
                  "dali_cpu pipeline is not compatible with multi_gpu mode :<"
              )
  
          def dali_dataset_fn(input_context):
              with tf.device(f"/gpu:{input_context.input_pipeline_id}"):
                  device_id = input_context.input_pipeline_id
                  num_shards = input_context.num_input_pipelines
                  return EfficientDetPipeline(
                      params,
                      int(total_batch_size / num_shards),
                      args,
                      is_training=is_training,
                      num_shards=num_shards,
                      device_id=device_id,
                  ).get_dataset()
  
          input_options = tf.distribute.InputOptions(
              experimental_place_dataset_on_device=True,
              experimental_fetch_to_device=False,
              experimental_replication_mode=tf.distribute.InputReplicationMode.PER_REPLICA,
          )
  
          dataset = strategy.distribute_datasets_from_function(
              dali_dataset_fn, input_options
          )
  
      else:
          from pipeline.dali.efficientdet_pipeline import (
              EfficientDetPipeline,
          )
  
          cpu_only = pipeline == PipelineType.dali_cpu
          device = "/cpu:0" if cpu_only else "/gpu:0"
          with tf.device(device):
              dataset = EfficientDetPipeline(
                  params,
                  total_batch_size,
                  args,
                  is_training=is_training,
                  cpu_only=cpu_only,
              ).get_dataset()
  
      return dataset
  
imports: from typing import Text, Tuple, Union
  import tensorflow
  import argparse
  import tensorflow.compat.v1
  from enum import Enum
  from collections import namedtuple
  from tensorflow.python.eager import tape
  from tensorflow.python.tpu import tpu_function
label_desc: Use image node_id 3 for calling a global function with example usage: get_dataset(args, total_batch_size, is_training, params, strategy) and returns: dataset, EfficientDetPipeline