Ошибка при векторизации отображения с помощью tf.data.Dataset ()

Question

У меня есть набор данных изображения, который я получил с помощью tf.data.Dataset.list_files().

В моей функции .map() я читаю и декодирую изображения, как показано ниже:

def map_function(filepath):
    image = tf.io.read_file(filename=filepath)
    image = tf.image.decode_jpeg(image, channels=3)
    image = tf.image.convert_image_dtype(image, tf.float32)
    image = tf.image.resize(image, [IMAGE_WIDTH, IMAGE_HEIGHT])
    return image

Если Я использую (это работает ниже)

 dataset = tf.data.Dataset.list_files(file_pattern=...)
 dataset = dataset.map(map_function)
 for image in dataset.as_numpy_iterator():
    #Correctly outputs the numpy array, no error is displayed/encountered
    print(image)

Однако, если я использую (это ниже, выдает ошибку):

  dataset = tf.data.Dataset.list_files(file_pattern=...)
  dataset = dataset.batch(32).map(map_function)
  for image in dataset.as_numpy_iterator():
    #Error is displayed 
      print(image)

ValueError: Shape должен иметь ранг 0, но это ранг 1 для 'ReadFile' (op: 'ReadFile') с формами ввода: [?].

Теперь, согласно этому: https://www.tensorflow.org/guide/data_performance#vectorizing_mapping, код не должен завершиться ошибкой и шаг предварительной обработки должен быть оптимизирован (пакетная обработка по сравнению с одноразовой обработкой).

Где ошибка в моем коде?

*** Если я использую map().batch(), он работает нормально