def METHOD_NAME(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