Image semantic segmentation example with tf.data object on multi gpu training

Hello, I am training image semantic segmentation network on multiple gpu (4 gpus). Currently my data pipeline using tf.data.experimental_distribute_dataset from tensorflow as data pipeline and mirrored strategy. I want to use DALI library for faster performance. I have checkout some examples, but non of them helped me. Below is the code that I have added but not working. Any suggestion would be helpful.

Directory Structure:

|-- Root
|  |-- train
|  |  |-- images (png images)
|  |  |-- labels   (png label segmentation mask)
|  |-- valid
|  |  |-- images
|  |  |-- labels
|  |-- test
|  |  |-- images
|  |  |-- labels

I have added code below in my example.

@pipeline_def(batch_size=args.batch_size)
def data_pipeline(shard_id):
    pngs = fn.readers.caffe2(
        path=os.path.join(args.dataset, "train", "images"), random_shuffle=True, shard_id=shard_id, num_shards=args.num_gpus)
    labels = fn.readers.caffe2(
        path=os.path.join(args.dataset, "train", "labels"), random_shuffle=True, shard_id=shard_id, num_shards=args.num_gpus)

    images = fn.decoders.image(pngs, device='mixed', output_type=types.RGB)
    labels = fn.decoders.image(labels, device='mixed', output_type=types.RGB)

    # images = fn.crop_mirror_normalize(
    #     images, dtype=types.FLOAT, std=[255.], output_layout="CHW")

   ???????? How can I apply crop without normalize on labels ????????
   
    return images, labels


shapes = (
    (args.batch_size, args.crop_height, args.crop_width),
    (args.batch_size))

dtypes = (
    tf.float32,
    tf.int32)


def dataset_fn(input_context):
    with tf.device("/gpu:{}".format(input_context.input_pipeline_id)):
        device_id = input_context.input_pipeline_id
        return dali_tf.DALIDataset(
            pipeline=data_pipeline(
                device_id=device_id, shard_id=device_id),
            batch_size=args.batch_size,
            output_shapes=shapes,
            output_dtypes=dtypes,
            device_id=device_id)


input_options = tf.distribute.InputOptions(
    experimental_place_dataset_on_device=True,
    experimental_fetch_to_device=False,
    experimental_replication_mode=tf.distribute.InputReplicationMode.PER_REPLICA)



strategy = tf.distribute.MirroredStrategy(
    [device.name[-5:] for device in tf.config.experimental.list_physical_devices('GPU')][:args.num_gpus])

train_dataset = strategy.distribute_datasets_from_function(dataset_fn, input_options)
val_batches = strategy.distribute_datasets_from_function(dataset_fn, input_options)

Please let me know If any additional information is needed. Thank you