U-Net dimension problem

All exercise tests have been passed… but I’m facing problem in training the model having following lines of codes -
EPOCHS = 40
VAL_SUBSPLITS = 5
BUFFER_SIZE = 500
BATCH_SIZE = 32
processed_image_ds.batch(BATCH_SIZE)
train_dataset = processed_image_ds.cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
train_dataset
print(processed_image_ds.element_spec)
model_history = unet.fit(train_dataset, epochs=EPOCHS)

Error:
(TensorSpec(shape=(96, 128, 3), dtype=tf.float32, name=None), TensorSpec(shape=(96, 128, 3), dtype=tf.float32, name=None))
Epoch 1/40

InvalidArgumentError

logits and labels must have the same first dimension, got logits shape [393216,23] and labels shape [1179648]
[[node sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits (defined at :9) ]] [Op:__inference_train_function_8668]

Function call stack:
train_function

I am unable to understand where exactly is my code wrong.

Hey,

I think, there should to be an issue with the dimensions of the second tensor of the Error message. It should have one channel since it is the mask, so the dimensions should be (96, 128, 3). I would suggest you checking your code and finding if anywhere you accidentaly didn’t modify the dimensions for the masks. It could be in the block 2.2 Preprocess your data..
Hope this helps.

Henrikh