Erro de formato de entrada ao atualizar CNN pré-treinado de classificação binária para multiclassificação

Tenho um conjunto de dados de 3 classes de imagens, subdivididas em pastas de treinamento/validação/teste:

new_base_dir = '/Users/.../img_dir'
import os
from tensorflow.keras.utils import image_dataset_from_directory
  
train_dataset = image_dataset_from_directory(
    os.path.join(new_base_dir, 'train'),
    image_size=(180, 180),
    batch_size=10)
validation_dataset = image_dataset_from_directory(
    os.path.join(new_base_dir, 'validation'),
    image_size=(180, 180),
    batch_size=10)
test_dataset = image_dataset_from_directory(
    os.path.join(new_base_dir, 'test'),
    image_size=(180, 180),
    batch_size=10)

Eu verifiquei que existem 3 classes:

train_dataset
Found 3000 files belonging to 3 classes
<BatchDataset element_spec=(TensorSpec(shape=(None, 180, 180, 3), dtype=tf.float32, name=None), TensorSpec(shape=(None,), dtype=tf.int32, name=None))>

Estou primeiro carregando as imagens em um Xceptionmodelo pré-treinado:

# load a pre-trained model
import keras
from keras import layers
from keras.applications import VGG16

conv_base = keras.applications.xception.Xception(
    weights="imagenet",
    include_top=False,
    input_shape=(180, 180, 3))

import numpy as np
  
def get_features_and_labels(dataset):
    all_features = []
    all_labels = []
    for images, labels in dataset:
        preprocessed_images = keras.applications.xception.preprocess_input(images)
        features = conv_base.predict(preprocessed_images)
        all_features.append(features)
        all_labels.append(labels)
    return np.concatenate(all_features), np.concatenate(all_labels)
  
train_features, train_labels =  get_features_and_labels(train_dataset)
val_features, val_labels =  get_features_and_labels(validation_dataset)
test_features, test_labels =  get_features_and_labels(test_dataset)

No entanto, quando coloco em camadas um classificador densamente conectado, recebo um erro:

inputs = keras.Input(shape=(6, 6, 2048))
x = layers.Flatten()(inputs)             
x = layers.Dense(256)(x)
x = layers.Dropout(0.5)(x)
# the commented lines work fine for a binary classification model
#outputs = layers.Dense(1, activation="sigmoid")(x)
#model.compile(loss="binary_crossentropy",
#              optimizer="rmsprop",
#              metrics=["accuracy"])
outputs = layers.Dense(3, activation="softmax")(x)
model = keras.Model(inputs, outputs)
model.compile(loss="categorical_crossentropy",
              optimizer="rmsprop",
              metrics=["accuracy"])

Aqui está o erro:

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