Модель Tensorflow генерирует ошибку недопустимого аргумента

Question

Я новичок в TensorFlow и работаю над проблемой классификации изображений. Прежде всего, мне пришлось добавить метки к изображениям с помощью класса ImageDataGenerator. Вот код, который я использовал для этого.

class_names = ['potholes', 'normal']
BATCH_SIZE = 32
IMG_HEIGHT = 224
IMG_WIDTH = 224

train_dir = r'/content/drive/My Drive/Colab Notebooks/Pothole Dataset/train'
train_data_gen = image_generator.flow_from_directory(directory=train_dir,
                                                     batch_size=BATCH_SIZE,
                                                     shuffle=True,
                                                     target_size=(IMG_HEIGHT, IMG_WIDTH),
                                                     classes=class_names)


val_dir =  r'/content/drive/My Drive/Colab Notebooks/Pothole Dataset/val'
val_data_gen = image_generator.flow_from_directory(directory=val_dir,
                                                     batch_size=BATCH_SIZE,
                                                     shuffle=True,
                                                     target_size=(IMG_HEIGHT, IMG_WIDTH),
                                                     classes=class_names)

После этого я использовал трансферное обучение для создания модели с использованием предварительно обученной модели mobil enet. Код ниже:

URL = "https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4"
feature_extractor = hub.KerasLayer(URL, input_shape=(IMG_WIDTH, IMG_HEIGHT, 3))

feature_extractor.trainable = False

model = tf.keras.Sequential([
  feature_extractor,
  tf.keras.layers.Dense(2)
])

model.summary()

model.compile(
  optimizer='adam',
  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
  metrics=['accuracy'])

Когда я вызвал метод fit_generator, я получил следующую ошибку:

InvalidArgumentError:  logits and labels must have the same first dimension, got logits shape [32,2] and labels shape [64]
     [[node sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits (defined at <ipython-input-83-86de6a95e60e>:8) ]] [Op:__inference_train_function_18463]

Function call stack:
train_function

Пожалуйста, я хотел бы знать причину этой ошибки и как ее отладить это. Спасибо.

PS Вот вызов метода fit_generator:

EPOCHS = 6

history = model.fit_generator(train_data_gen,
                              steps_per_epoch=int(np.ceil(no_of_training_images/float(BATCH_SIZE))),
                              epochs=EPOCHS,
                              validation_data=val_data_gen,
                              validation_steps=int(np.ceil(no_of_val_images/float(BATCH_SIZE)))
                              )