Keras model.fit() returns large array of int with node sparse_categorical_crossentropy error

Question

I am using the example Keras provides to attempt multiclass semantic segmentation using DeepLabV3+ (https://keras.io/examples/vision/deeplabv3_plus/) on images of robotic surgery to segment the surgical instrument and its different parts (shaft, joint, tip). To fit my needs I changed the IMAGE_SIZE = 1024 and NUM_CLASSES = 4 (shaft, joint, tip, background), NUM_TRAIN = 216, NUM_VAL = 54.

When I compile and train the model:

loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.001),
    loss=loss,
    metrics=["accuracy"]
)

history = model.fit(train_dataset, validation_data=val_dataset, epochs=25)

The output is: Epoch 1/25 "--- A large array of integers ----" [[{{node sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits}}]] [Op:__inference_train_function_165277]

The error I receive is:

Invalid Argument Error: Graph Execution error: Detected at node sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits

I am a beginner to python/computer vision, and any help would be greatly appreciated. Thank you in advance!

Attempted changing:

1. changing from logits = FALSE in loss function
2. changing input image size back to 512 (as in the keras example)