Как загрузить мой собственный набор данных, чтобы создать DCGAN в тензорном потоке?

Я пытался следовать руководству здесь, чтобы создать DCGAN, изменив только мои входные данные и Мне не удалось загрузить данные в надлежащую форму для обучения моей модели. Вот мой код:

import pathlib
data_dir = pathlib.Path(data_dir)

list_ds = tf.data.Dataset.list_files(str(data_dir/'*/*'))

def decode_img(img):
  # convert the compressed string to a 3D uint8 tensor
  img = tf.image.decode_jpeg(img, channels=1)
  # Use `convert_image_dtype` to convert to floats in the [0,1] range.
  img = tf.image.convert_image_dtype(img, tf.float32)
  # resize the image to the desired size.
  return tf.image.resize(img, [IMG_WIDTH, IMG_HEIGHT])
def process_path(file_path):
  # load the raw data from the file as a string
  img = tf.io.read_file(file_path)
  img = decode_img(img)
  return img
labeled_ds = list_ds.map(process_path)

и оттуда я в значительной степени скопировал руководство, изменив параметры для соответствия моему набору данных, и получил такой конечный результат:

train(labeled_ds, EPOCHS)

но при этом получаю следующую ошибку.

ValueError: Input 0 of layer sequential_1 is incompatible with the layer: expected ndim=4, found ndim=3. Full shape received: [224, 320, 1]

Как я могу добавить дополнительное измерение в набор данных тензорного потока? Или, скорее, я что-то делаю неправильно, а не то, как набор данных MNIST был загружен в учебнике? Спасибо