Issue with Keras input array when implementing Convolutional Neural Network - machine-learning

Question

I'm trying to implement a convolutional neural network within Keras using a TF backend for image segmentation of 111 images of size 141 x 166. When I run the code below, I get the error message:
Error when checking target: expected dense_36 to have 2 dimensions, but got array with shape (88, 141, 166, 1)
My X_train variable is the shape (88, 141, 166, 1) as well as the y_train variable. My X_test variable is the shape (23, 141, 166, 1) as well as the y_test variable, as split by the function train_test_split from sklearn.
I'm not sure what the error message means as per dense_36. I have tried using the Flatten() function before fitting the model, but it says that I have a ndim = 2 and cannot be flattened.
# set input
batch_size = 111
num_epochs = 50
img_rows = 141
img_cols = 166
input_shape = (img_rows, img_cols, 1)
num_classes = img_rows*img_cols
# split training and test data
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.2, random_state = 4)
X_train = X_train.astype('float32')
X_test = X_train.astype('float32')
# CNN itself
model = Sequential()
model.add(Conv2D(32, kernel_size=(3,3), activation='relu',
input_shape=input_shape))
model.add(Conv2D(64, (3,3), activation='relu'))
model.add(MaxPooling2D(pool_size=(2,2)))
model.add(Dropout(0.5))
model.add(Flatten())
model.add(Dense(128, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(num_classes, activation='softmax'))
# compile CNN
model.compile(loss=keras.losses.categorical_crossentropy,
optimizer=keras.optimizers.Adadelta(), metrics=['accuracy'])
# fit CNN
model.fit(X_train, y_train, batch_size=batch_size, epochs=num_epochs,
verbose=1, validation_data=(X_test, y_test))
My model summary is:
Layer (type) Output Shape Param #
=================================================================
conv2d_35 (Conv2D) (None, 139, 164, 32) 320
_________________________________________________________________
conv2d_36 (Conv2D) (None, 137, 162, 64) 18496
_________________________________________________________________
max_pooling2d_18 (MaxPooling (None, 68, 81, 64) 0
_________________________________________________________________
dropout_35 (Dropout) (None, 68, 81, 64) 0
_________________________________________________________________
flatten_28 (Flatten) (None, 352512) 0
_________________________________________________________________
dense_33 (Dense) (None, 128) 45121664
_________________________________________________________________
dropout_36 (Dropout) (None, 128) 0
_________________________________________________________________
dense_34 (Dense) (None, 2) 258
_________________________________________________________________
Total params: 45,140,738
Trainable params: 45,140,738
Non-trainable params: 0
_________________________________________________________________
None