I’m trying to pass some inputs through a couple of layers, but it seems i’m misaligning my input shape somewhere such that it cannot pass properly through the layers thereby raising a value error as folllows:

(ds_train, ds_test), ds_info = tfds.load(
    "mnist",
    split=["train", "test"],
    shuffle_files=True,
    as_supervised=True,
    with_info=True
)

def normalize_img(image, label):
    return tf.cast(image, tf.float32)/255.0, label

AUTOTUNE = tf.data.experimental.AUTOTUNE
BATCH_SIZE = 64

ds_train = ds_train.map(normalize_img, num_parallel_calls=AUTOTUNE)
ds_train = ds_train.cache()
ds_train = ds_train.shuffle(ds_info.splits["train"].num_examples)
ds_train = ds_train.batch(BATCH_SIZE)
ds_train = ds_train.prefetch(AUTOTUNE)

ds_test = ds_test.map(normalize_img, num_parallel_calls=AUTOTUNE)
de_test = ds_test.batch(128)
ds_test = ds_test.prefetch(AUTOTUNE)

model = keras.Sequential ([
    keras.Input(shape=[28, 28, 1],),
    layers.Conv2D(32, 3, activation='relu', padding="same"),
    layers.Flatten(),
    layers.Dense(10),

])

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

)

model.fit(ds_train, epochs=5, verbose=2)
model.evaluate(ds_test)

ValueError: Exception encountered when calling Sequential.call().

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

Arguments received by Sequential.call():
  • inputs=tf.Tensor(shape=(28, 28, 1), dtype=float32)
  • training=False
  • mask=None

please help out