code: def get_dataset(
      device="cpu",
      device_id=0,
      shard_id=0,
      num_shards=1,
      fail_on_device_mismatch=True,
  ):
      pipeline = mnist_pipeline(
          4, data_path, device, device_id, shard_id, num_shards
      )
      shapes = ((BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE), (BATCH_SIZE,))
      dtypes = (tf.float32, tf.int32)
  
      daliset = dali_tf.DALIDataset(
          pipeline=pipeline,
          batch_size=BATCH_SIZE,
          output_shapes=shapes,
          output_dtypes=dtypes,
          num_threads=4,
          device_id=device_id,
          fail_on_device_mismatch=fail_on_device_mismatch,
      )
      return daliset
  
imports: import nvidia.dali.plugin.tf
  import os
  from nvidia.dali.pipeline import Pipeline
  import nvidia.dali.fn
  import nvidia.dali.types
  from test_utils_tensorflow import num_available_gpus
  from shutil import rmtree
  import tensorflow
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling a global function with example usage: get_dataset(device, device_id, shard_id, num_shards, fail_on_device_mismatch) and returns: daliset