MNIST puzzling error

I’m a noob still but am running into a puzzling error. I am playing with the MNIST tutorial. Here is the network:

        model = tf.keras.models.Sequential( [
        tf.keras.layers.Flatten(input_shape=(28, 28)),
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dense(10)
        ] )
        
       model.compile(
            optimizer=tf.keras.optimizers.Adam(0.001),
            loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
            metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],
        )

After training, I want to pass a single data point through the network so I do this:

mnist = tfds.load('mnist', split='train', as_supervised=True)
data_iter = mnist.repeat().as_numpy_iterator()
image_in = next( data_iter )
model( np.squeeze( image_in ) ) 

However, it appears the sqeeze operation in the first layer isn’t doing it’s job according to this error:

ValueError: Exception encountered when calling layer 'sequential_2' (type Sequential).

Input 0 of layer "dense_4" is incompatible with the layer: expected axis -1 of input shape to have value 784, but received input with shape (28, 28)

Can someone help me understand what’s going on here and what the appropriate way is to pass a data point through the model is?

Your help is greatly appreciated.

I found my issue: the input needs to be 1xMxM and not MxMx1 (which is returned by the iterator).