TensorFlow GradCAM — model.fit () — ValueError: формы (нет, 1) и (нет, 2) несовместимы

В рамках задания 4 курса Coursera CV TF мой код не работает model.fit()

model.compile(loss='categorical_crossentropy',metrics=
   ['accuracy'],optimizer=tf.keras.optimizers.RMSprop(lr=0.001))
# shuffle and create batches before training

model.fit(train_batches,epochs=25)

с ошибкой:

ValueError: Shapes (None, 1) and (None, 2) are incompatible

Любой намек на то, откуда может возникнуть проблема? Я подозреваю неправильный формат или тип для train_batches:

train_data = tfds.load('cats_vs_dogs', split='train[:80%]', as_supervised=True) 
augmented_training_data = train_data.map(augmentimages)
train_batches = augmented_training_data.batch(32)