code: def get_dataset_multi_gpu(strategy):
      def dataset_fn(input_context):
          with tf.device(
              "/gpu:{}".format(input_context.input_pipeline_id)
          ):
              device_id = input_context.input_pipeline_id
              return get_dataset(
                  "gpu", device_id, device_id, num_available_gpus()
              )
  
      input_options = tf.distribute.InputOptions(
          experimental_place_dataset_on_device=True,
          experimental_fetch_to_device=False,
          experimental_replication_mode=tf.distribute.InputReplicationMode.PER_REPLICA,
      )
  
      train_dataset = strategy.distribute_datasets_from_function(
          dataset_fn, input_options
      )
      return train_dataset
  
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 3 for calling a global function with example usage: get_dataset_multi_gpu(strategy) and returns: train_dataset, get_dataset