Метод .map () набора данных TensorFlow не работает для встроенных функций tf.keras.preprocessing.image

Question

Я загружаю набор данных как таковой:

import tensorflow_datasets as tfds

ds = tfds.load(
    'caltech_birds2010',
    split='train',
    as_supervised=False)

И эта функция отлично работает:

import tensorflow as tf

@tf.function
def pad(image,label):
    return (tf.image.resize_with_pad(image,32,32),label)

ds = ds.map(pad)

Но когда я пытаюсь сопоставить другую встроенную функцию

from tf.keras.preprocessing.image import random_rotation

@tf.function
def rotate(image,label):
    return (random_rotation(image,90), label)

ds = ds.map(rotate)

Я получаю следующую ошибку:

AttributeError: объект 'Tensor' не имеет атрибута 'ndim'

Это не единственная функция, которая дает мне проблемы, и это происходит с декоратором @tf.function или без него.

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