Применение стилизации к партиям вызывает проблему

Question

В настоящее время я использую модуль стилизации TF-Hub для серии экспериментов, в которых я применяю стилизацию случайным образом к набору изображений.

После применения стилизации изображения выглядят следующим образом:

Я использую Cats-vs.-Dogs набор данных, и я загружаю его, используя следующий листинг кода:

train_ds, validation_ds = tfds.load(
    'cats_vs_dogs',
    split=['train[:85%]', 'train[85%:]'],
    as_supervised=True
)

После этого я создаю служебную функцию для применения модуля стилизации к данному изображению, и я отображаю эту служебную программу следующим образом -

train_ds = (
    train_ds
    .map(preprocess_image, num_parallel_calls=AUTO)
    .shuffle(1024)
    .map(apply_random_stylization, num_parallel_calls=AUTO)
    .batch(BATCH_SIZE)
    .prefetch(AUTO)
)

apply_random_stylization определяется как -

def apply_random_stylization(image, label):
    image = random_apply(stylize_image, image, p=0.8)
    return (image, label)

def random_apply(func, x, p):
        return tf.cond(
          tf.less(tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32),
                  tf.cast(p, tf.float32)),
          lambda: func(x),
          lambda: x)

def stylize_image(image):
    stylized_image = HUB_MODULE(tf.expand_dims(image, 0), tf.constant(STYLE_IMG))[0]
    stylized_image = tf.squeeze(stylized_image)
    return tf.clip_by_value(stylized_image, 0., 1.)

Вот Colab Gist , воспроизводящий эту проблему.

Любая помощь приветствуется. Спасибо!