shapes are incompatible error in CNN model using tensorflow

Question

I'm try to make a convolutional network for roof top classification with two classes 0,1.

The images size are 128x128, I use 3 more convolutional layer to reduce it to 6 x 5 as adviced here

I'm very new at convolutional network. So, I may be wrong in some place, so feel free to advice me on this.

def cnn_model_fn(features,labels,mode):
    input_layer = tf.reshape(features, [-1,128,128,3])

    #Convolution layer #1
    conv1 = tf.layers.conv2d(input_layer,filters=28,kernel_size=[5,5],padding="same",activation=tf.nn.relu)

    #Pooling layer #1
    pool1=tf.layers.max_pooling2d(inputs=conv1,pool_size=[2,2],strides=2)

    #Convolution layer #2
    conv2 = tf.layers.conv2d(pool1,filters=14,kernel_size=[5,5],padding="same",activation=tf.nn.relu)

    #Pooling layer #2
    pool2=tf.layers.max_pooling2d(inputs=conv2,pool_size=[2,2],strides=2)

    # Dense Layer
    pool2_flat = tf.reshape(pool2, [-1, 128 * 128 * 3])
    dense = tf.layers.dense(inputs=pool2_flat, units=128, activation=tf.nn.relu)
    dense2 = tf.layers.dense(inputs=dense, units=64, activation=tf.nn.relu)
    dense3 = tf.layers.dense(inputs=dense2, units=32, activation=tf.nn.relu)
    dropout = tf.layers.dropout(inputs=dense3, rate=0.4, training=mode == learn.ModeKeys.TRAIN)

    # Logits Layer
    logits = tf.layers.dense(inputs=dropout, units=2)

    loss = None
    train_op = None

    # Calculate Loss (for both TRAIN and EVAL modes)
    if mode != learn.ModeKeys.INFER:
        onehot_labels = tf.one_hot(indices=tf.cast(labels, tf.int32), depth=4)
        loss = tf.losses.softmax_cross_entropy(
        onehot_labels=onehot_labels, logits=logits)

    # Configure the Training Op (for TRAIN mode)
    if mode == learn.ModeKeys.TRAIN:
        train_op = tf.contrib.layers.optimize_loss(
        loss=loss,
        global_step=tf.contrib.framework.get_global_step(),
        learning_rate=0.001,
        optimizer="SGD")

    # Generate Predictions
    predictions = {
      "classes": tf.argmax(
          input=logits, axis=1),
      "probabilities": tf.nn.softmax(
          logits, name="softmax_tensor")
      }

#Image data into train and test

data_dir='C:/Users/user/Desktop/exzeo/Category1/' 
data_dir2='C:/Users/user/Desktop/exzeo/Category2/'
data_images=[]
data_labels=[]
for image_path in glob.glob(data_dir+"*.png"):
    data_images.append(cv2.resize(cv2.imread(image_path),(128,128)))
    data_labels.append(np.array(0))
for image_path1 in glob.glob(data_dir2+"*.png"):
    data_images.append(cv2.resize(cv2.imread(image_path1),(128,128)))
    data_labels.append(np.array(1))
data=np.float32(np.asarray(data_images))
labels=np.float32(np.asarray(data_labels))
train_data, eval_data, train_labels, eval_labels = train_test_split(data, labels, test_size=0.20, random_state=42)

Model Fitting

def main(unused_argv):
    building_classifier = learn.Estimator(model_fn=cnn_model_fn)
    #     # Set up logging for predictions
    tensors_to_log = {"probabilities": "softmax_tensor"}
    logging_hook = tf.train.LoggingTensorHook(tensors=tensors_to_log, every_n_iter=50)
    building_classifier.fit(x=train_data,y=train_labels,batch_size=128,steps=10000,monitors=[logging_hook])
    #     # Configure the accuracy metric for evaluation
    metrics = {"accuracy":learn.MetricSpec(metric_fn=tf.metrics.accuracy, prediction_key="classes"),}
    eval_results = building_classifier.evaluate(x=eval_data, y=eval_labels, metrics=metrics)
    print(eval_results)
    if __name__ == "__main__":
        tf.app.run()

As, it throw me a error as ValueError: Shapes (?, 2) and (?, 4) are incompatible