ValueError ao adicionar tf.keras.metrics.Precision para modelar métricas no TensorFlow

Estou tentando aplicar a aprendizagem por transferência usando MobileNet em um conjunto de dados personalizado. Meu código funciona bem até eu adicionar tf.keras.metrics.Precision(name="precision") às métricas do modelo. Depois de adicionar esta métrica, encontro o seguinte erro durante model.fit:

update_confusion_matrix_variables y_pred.shape.assert_is_compatible_with(y_true.shape) ValueError: Shapes (None, 4) and (None, 1) are incompatible

O bloco de código relevante:

import tensorflow as tf
from tensorflow.keras.utils import image_dataset_from_directory as get_dataset

img_size = (224, 224)
preprocessing_layer = tf.keras.applications.mobilenet.preprocess_input
img_shape = img_size + (3,)

base_model = tf.keras.applications.MobileNet(
    input_shape=img_shape,
    include_top=False,
    weights='imagenet',
)

batch_size = 32
train_set = get_dataset(
    'some path',
    shuffle=True,
    batch_size=batch_size,
    image_size=img_size,
)

class_names = train_set.class_names
num_classes = len(class_names)

base_model.trainable = False
global_average_layer = tf.keras.layers.GlobalAveragePooling2D()

inputs = tf.keras.Input(shape=img_shape)
k = preprocessing_layer(inputs)
k = base_model(k, training=False)
k = global_average_layer(k)
k = tf.keras.layers.Dropout(0.2)(k)

prediction_layer = tf.keras.layers.Dense(
    num_classes,
    activation="softmax"
)
outputs = prediction_layer(k)

model = tf.keras.Model(inputs, outputs)

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(),
    metrics=[
        tf.keras.metrics.SparseCategoricalAccuracy(name='accuracy'),
        tf.keras.metrics.Precision(name="precision"),
    ]
)

AUTOTUNE = tf.data.AUTOTUNE

train_set = train_set.prefetch(buffer_size=AUTOTUNE)

epochs = 1
history = model.fit(
    train_set,
    epochs=epochs,
)

O conjunto de dados é organizado em subdiretórios onde cada nome de subdiretório é o rótulo da classe. Existem quatro rótulos de classe.

O carregamento e o pré-processamento do conjunto de dados parecem funcionar corretamente, pois o modelo treina sem a tf.keras.metrics.Precisionmétrica.

O problema surge especificamente ao adicionar a métrica de precisão.

epochsé intencionalmente definido 1para fins experimentais.