Keep Getting ValueError : Shapes (None, None) and (None, None, None, 1174) are incompatible when fitting UNet

Question

I'm building my first u-net, and having a bit of trouble matching the output to the labels.

My code is as follows:

import tensorflow as tf
from keras import layers, models
from keras.preprocessing.image import ImageDataGenerator

def unet(input_shape=(160, 160, 3), num_classes=1174):
    inputs = tf.keras.Input(shape=input_shape)

    # Encoder (contracting path)
    conv1 = layers.Conv2D(64, 3, activation='relu', padding='same')(inputs)
    conv1 = layers.Conv2D(64, 3, activation='relu', padding='same')(conv1)
    pool1 = layers.MaxPooling2D(pool_size=(2, 2))(conv1)

    conv2 = layers.Conv2D(128, 3, activation='relu', padding='same')(pool1)
    conv2 = layers.Conv2D(128, 3, activation='relu', padding='same')(conv2)
    pool2 = layers.MaxPooling2D(pool_size=(2, 2))(conv2)

    # Bottleneck
    conv3 = layers.Conv2D(256, 3, activation='relu', padding='same')(pool2)
    conv3 = layers.Conv2D(256, 3, activation='relu', padding='same')(conv3)

    # Decoder (expansive path)
    up4 = layers.UpSampling2D(size=(2, 2))(conv3)
    up4 = layers.Conv2D(128, 2, activation='relu', padding='same')(up4)
    concat4 = layers.Concatenate()([conv2, up4])
    conv4 = layers.Conv2D(128, 3, activation='relu', padding='same')(concat4)
    conv4 = layers.Conv2D(128, 3, activation='relu', padding='same')(conv4)

    up5 = layers.UpSampling2D(size=(2, 2))(conv4)
    up5 = layers.Conv2D(64, 2, activation='relu', padding='same')(up5)
    concat5 = layers.Concatenate()([conv1, up5])
    conv5 = layers.Conv2D(64, 3, activation='relu', padding='same')(concat5)
    conv5 = layers.Conv2D(64, 3, activation='relu', padding='same')(conv5)

    # Output layer
    outputs = layers.Conv2D(num_classes, 1, activation='softmax')(conv5)

    # Define the model
    model = models.Model(inputs=inputs, outputs=outputs)
    return model

train_directory = '/mnt/c/Users/user1/my_repo/data_folder/train'
test_directory = '/mnt/c/Users/user1/my_repo/data_folder/test'

# Data generators
train_datagen = ImageDataGenerator(rescale=1./255)
test_datagen = ImageDataGenerator(rescale=1./255)

train_generator = train_datagen.flow_from_directory(
    train_directory,
    target_size=(160, 160),
    batch_size=8,
    class_mode='categorical'
)

test_generator = test_datagen.flow_from_directory(
    test_directory,
    target_size=(160, 160),
    batch_size=8,
    class_mode='categorical'
)

#Fit the mode;

model = unet()
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

model.fit(
    train_generator,
    steps_per_epoch=len(train_generator),
    epochs=10,
    validation_data=test_generator,
    validation_steps=len(test_generator)
)

It gives me the error:

File "/home/user/.local/lib/python3.10/site-packages/keras/src/backend.py", line 5573, in categorical_crossentropy
        target.shape.assert_is_compatible_with(output.shape)

ValueError: Shapes (None, None) and (None, None, None, 1174) are incompatible

I have tried a variety of things, which ended up giving me different errors. I tried changing the loss function as well as the data generators to sparse_categorical_crossentropy, and spare respectively, but that resulted in another mismatching error where values were not broadcastable.