I am trying to train a simple CNN on custom image data. While training with raw numpy data, the model learns as expected. However when I convert the same raw numpy data to dataset using tf.data.Dataset.from_tensor_slices, the training goes completely goes haywire, and the model does not learn anything.

The relevant part of code look like this

    train_samples_arr = np.empty((num_train_samples, 64, 64, 1), dtype=np.float64)
    train_labels_arr = np.empty((num_train_samples,), dtype=np.float64)
    count = 0
    for npz, label in train_samples:
        for cidx, img in enumerate(npz):
            train_samples_arr[count,] = img.reshape(64, 64, 1).astype("float64")/255.0
            train_labels_arr[count] = label
            count += 1

    #Creating tf.data.Dataset here
    BATCH_SIZE = 32
    train_dataset = tf.data.Dataset.from_tensor_slices((train_samples_arr, train_labels_arr))
    test_dataset = tf.data.Dataset.from_tensor_slices((test_samples_arr, test_labels_arr))
    train_dataset = train_dataset.batch(BATCH_SIZE)
    test_dataset = test_dataset.batch(BATCH_SIZE)

    #1. With tf.data.Dataset.from_tensor_slices
    model.fit(train_dataset, validation_data = test_dataset, epochs=3)

    #2. Seperate run with raw numpy
    model.fit(train_samples_arr, train_labels_arr, batch_size=32, validation_data = (test_samples_arr, test_labels_arr), epochs=3)

The training outputs look very different.

1. With tf datasets:

782/782 ━━━━━━━━━━━━━━━━━━━━ 21s 19ms/step - loss: 3.0736 - sparse_categorical_accuracy: 0.7187 - 
val_loss: 7.2740 - val_sparse_categorical_accuracy: 0.0299
Epoch 2/3
782/782 ━━━━━━━━━━━━━━━━━━━━ 14s 18ms/step - loss: 2.8841 - sparse_categorical_accuracy: 0.1824 - 
val_loss: 11.5464 - val_sparse_categorical_accuracy: 0.0299
Epoch 3/3
782/782 ━━━━━━━━━━━━━━━━━━━━ 14s 18ms/step - loss: 2.6524 - sparse_categorical_accuracy: 0.4601 - 
val_loss: 3.2222 - val_sparse_categorical_accuracy: 0.0410

2. With Raw numpy

782/782 ━━━━━━━━━━━━━━━━━━━━ 20s 19ms/step - loss: 2.1845 - sparse_categorical_accuracy: 0.3749 - 
val_loss: 0.3580 - val_sparse_categorical_accuracy: 0.8877
Epoch 2/3
782/782 ━━━━━━━━━━━━━━━━━━━━ 14s 17ms/step - loss: 0.5643 - sparse_categorical_accuracy: 0.8209 - 
val_loss: 0.2546 - val_sparse_categorical_accuracy: 0.9204
Epoch 3/3
782/782 ━━━━━━━━━━━━━━━━━━━━ 14s 17ms/step - loss: 0.4270 - sparse_categorical_accuracy: 0.8684 - 
val_loss: 0.2536 - val_sparse_categorical_accuracy: 0.9200

Looking at val_sparse_categorical_accuracy obviously shows that the model when trained on datasets is not learning anything.

What am I doing wrong?