ValueError: Can not squeeze dim[1], expected a dimension of 1, got 24 CNN

Question

Good morning all, I want to use a cnn model which takes in an image in jpg and predict the (x,y) coordinates of the center of one of the structures in the image. My image has a shape of (512,512,3) and my coordinate labels is a 24 by 2 array. I use this cnn model

import tensorflow as tf

from tensorflow.keras import datasets, layers, models
import matplotlib.pyplot as plt
model = models.Sequential()
model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(512, 512, 3)))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.Flatten())
model.add(layers.Dense(64, activation='relu'))
model.add(layers.Dense(10))
model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
history = model.fit(images, labels, epochs=10, 
                    validation_data=(test_images, test_labels))

and it is giving me error below:

ValueError                                Traceback (most recent call last)
Cell In[34], line 5
      1 model.compile(optimizer='adam',
      2               loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
      3               metrics=['accuracy'])
----> 5 history = model.fit(images, labels, epochs=10, 
      6                     validation_data=(test_images, test_labels))

File ~/anaconda3/lib/python3.11/site-packages/keras/src/utils/traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     67     filtered_tb = _process_traceback_frames(e.__traceback__)
     68     # To get the full stack trace, call:
     69     # `tf.debugging.disable_traceback_filtering()`
---> 70     raise e.with_traceback(filtered_tb) from None
     71 finally:
     72     del filtered_tb

File /var/folders/_k/pyjyt41x19x50pc7dsq_4h_40000gn/T/__autograph_generated_fileb4jt3ym7.py:15, in outer_factory.<locals>.inner_factory.<locals>.tf__train_function(iterator)
     13 try:
     14     do_return = True
---> 15     retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)
     16 except:
     17     do_return = False

ValueError: in user code:
ValueError: Can not squeeze dim[1], expected a dimension of 1, got 24 for '{{node Squeeze}} = Squeeze[T=DT_INT64, squeeze_dims=[-1]](ArgMax)' with input shapes: [?,24].

Here are what my train_x and train_y looks like:

([array([[[0.00392157, 0.00392157, 0.00392157],
          [0.00392157, 0.00392157, 0.00392157],
          [0.00392157, 0.00392157, 0.00392157],
          ...,
          [0.06666667, 0.06666667, 0.06666667],
          [0.07058824, 0.07058824, 0.07058824],
          [0.08627451, 0.08627451, 0.08627451]],
  
         [[0.        , 0.        , 0.        ],
          [0.        , 0.        , 0.        ],
          [0.        , 0.        , 0.        ],
          ...,
          [0.08235294, 0.08235294, 0.08235294],
          [0.07450981, 0.07450981, 0.07450981],
          [0.07058824, 0.07058824, 0.07058824]],
  
         [[0.        , 0.        , 0.        ],
          [0.        , 0.        , 0.        ],
          [0.        , 0.        , 0.        ],
          ...,
          [0.07058824, 0.07058824, 0.07058824],
          [0.0627451 , 0.0627451 , 0.0627451 ],
          [0.05882353, 0.05882353, 0.05882353]],
  
         ...,
  
         [[0.09019608, 0.09019608, 0.09019608],
          [0.09411765, 0.09411765, 0.09411765],
          [0.09411765, 0.09411765, 0.09411765],
          ...,
          [0.        , 0.        , 0.        ],
          [0.        , 0.        , 0.        ],
          [0.        , 0.        , 0.        ]],
  
         [[0.08235294, 0.08235294, 0.08235294],
          [0.09019608, 0.09019608, 0.09019608],
          [0.09803922, 0.09803922, 0.09803922],
          ...,
          [0.        , 0.        , 0.        ],
          [0.        , 0.        , 0.        ],
          [0.        , 0.        , 0.        ]],
  
         [[0.08627451, 0.08627451, 0.08627451],
          [0.09411765, 0.09411765, 0.09411765],
          [0.09411765, 0.09411765, 0.09411765],
          ...,
          [0.        , 0.        , 0.        ],
          [0.        , 0.        , 0.        ],
          [0.        , 0.        , 0.        ]]], dtype=float32)],
 [array(['S1', 139.0, 575.0, 'L5', 128.0, 555.0, 'L4', 119.0, 522.0, 'L3',
         119.0, 490.0, 'L2', 127.0, 457.0, 'L1', 137.0, 431.0, 'T12', 150.0,
         399.0, 'T11', 162.0, 373.0, 'T10', 173.0, 351.0, 'T9', 181.0,
         327.0, 'T8', 188.0, 310.0, 'T7', 194.0, 288.0, 'T6', 197.0, 263.0,
         'T5', 196.0, 245.0, 'T4', 192.0, 228.0, 'T3', 187.0, 211.0, 'T2',
         177.0, 196.0, 'T1', 165.0, 179.0, 'C7', 158.0, 159.0, 'C6', 154.0,
         134.0, 'C5', 156.0, 119.0, 'C4', 154.0, 105.0, 'C3', 156.0, 88.0,
         'C2', 162.0, 71.0], dtype=object)])

Thank you so much!

I do not know whether cnn is the right model for the purpose I want to achieve or there are something wrong with my model design.