def METHOD_NAME(self):
    os.makedirs(self.download_directory, exist_ok=True)
    # Use a file lock so only one worker on each node does the download
    with filelock.FileLock(os.path.join(self.download_directory, "download.lock")):
        dataset = tfds.load(
            "oxford_iiit_pet:3.*.*",
            split="train",
            with_info=False,
            data_dir=self.download_directory,
        )
    def load_image_train(datapoint):
        input_image = tf.image.resize(datapoint["image"], (128, 128))
        input_mask = tf.image.resize(datapoint["segmentation_mask"], (128, 128))
        if np.random.uniform(()) > 0.5:
            input_image = tf.image.flip_left_right(input_image)
            input_mask = tf.image.flip_left_right(input_mask)
        input_image, input_mask = self.normalize(input_image, input_mask)
        return input_image, input_mask
    train = dataset.map(load_image_train, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    train = self.context.wrap_dataset(train)
    train_dataset = (
        train.cache()
        .shuffle(self.context.get_data_config().get("BUFFER_SIZE"))
        .batch(self.context.get_per_slot_batch_size())
        .repeat()
    )
    train_dataset = train_dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
    return train_dataset