[Python] Keras AUC metric breaking on a multi-class problem

I'm attempting to train a multi-class image classification model based on standardised 128x128 images. When I use "accuracy" alone as a metric, I have no issues. When I introduce AUC, I can't train the model.

model = tf.keras.Sequential([
tf.keras.layers.Rescaling(1./255),
tf.keras.layers.Conv2D(32, 3, activation='relu'),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Dropout(dropout_rate),
tf.keras.layers.Conv2D(32, 3, activation='relu'),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Dropout(dropout_rate),
tf.keras.layers.Conv2D(32, 3, activation='relu'),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Dropout(dropout_rate),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(num_classes, activation="softmax")
])

model.compile(
optimizer='adam',
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=['accuracy',keras.metrics.AUC(curve="ROC",multi_label=False,from_logits=True)])


This is the exception thrown by Keras:

> File "C:\Users\ikhakoo\AppData\Local\Programs\Python\Python312\Lib\site-packages\keras\src\metrics\metrics_utils.py", line 272, in _update_confusion_matrix_variables_optimized
> File "C:\Users\ikhakoo\AppData\Local\Programs\Python\Python312\Lib\site-packages\keras\src\ops\math.py", line 86, in segment_sum
> File "C:\Users\ikhakoo\AppData\Local\Programs\Python\Python312\Lib\site-packages\keras\src\backend\tensorflow\math.py", line 23, in segment_sum
>data.shape = [32] does not start with segment_ids.shape = [128]
[[{{node UnsortedSegmentSum}}]] [Op:__inference_multi_step_on_iterator_130104]


I am not sure what the AUC metric expects in terms of the shapes - I notice that the first number corresponds to the batch size, but both numbers change when I vary the batch size. What am I doing wrong?

Continue reading...