Image Classification project. NaN Loss while training.

[santiagorg2401]

Hello,
I am currently preparing for the certification exam and while doing the checklist map I am training a few models. While training a feature extraction model and a normal sequential model I am facing a training issue, which is that the loss value remains at NaN at all times. This is a sample epoch log.

Epoch 1/100
858/858 [==============================] - 6s 5ms/step - loss: nan - accuracy: 0.0410 - val_loss: nan - val_accuracy: 0.0462

I am using the Sign Language MNIST dataset which comes in two separate .csv files for training and testing. I preprocessed them using the following code.

import pandas as pd
import numpy as np
import tensorflow as tf

train_df = pd.read_csv("/content/sign-language-mnist/sign_mnist_train/sign_mnist_train.csv")
test_df = pd.read_csv("/content/sign-language-mnist/sign_mnist_test/sign_mnist_test.csv")

X_train = train_df.drop(["label"], axis=1)
y_train = train_df["label"].to_numpy()
X_test = test_df.drop(["label"], axis=1)
y_test = test_df["label"].to_numpy()

X_train = X_train.values.reshape(-1, 28, 28, 1)
X_test = X_test.values.reshape(-1, 28, 28, 1)

X_train = np.repeat(X_train, 3, -1)
X_test = np.repeat(X_test, 3, -1)

plt.imshow(X_train[0], cmap="gray")
plt.axis(False)

X_train = tf.data.Dataset.from_tensor_slices(X_train)
y_train = tf.data.Dataset.from_tensor_slices(y_train)
X_test = tf.data.Dataset.from_tensor_slices(X_test)
y_test = tf.data.Dataset.from_tensor_slices(y_test)

X_train = X_train.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)
y_train = y_train.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)
X_test = X_test.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)
y_test = y_test.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)

train = tf.data.Dataset.zip((X_train, y_train))
test = tf.data.Dataset.zip((X_test, y_test))

The image plots correctly and the shapes are:

X_train: (27455, 28, 28, 3)
y_train: (27455,)
X_test: (7172, 28, 28, 3)
y_test: (7172,)

The model code is the following.

input = layers.Input(shape=(28, 28, 3,), name="input_layer")
x = layers.Normalization()(input)
x = layers.Conv2D(10, 3, activation="relu")(x)
x = layers.Conv2D(10, 3, activation="relu")(x)
x = layers.MaxPool2D()(x)
x = layers.Conv2D(10, 3, activation="relu")(x)
x = layers.Conv2D(10, 3, activation="relu")(x)
x = layers.MaxPool2D()(x)
x = layers.Flatten()(x)
output = layers.Dense(24, activation="softmax")(x)

model_6 = tf.keras.Model(input, output, name="model_6")

tf.random.set_seed(SEED)
model_6.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(),
                optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
                metrics=["accuracy"])
model_6_history = model_6.fit(train,
                              validation_data=test,
                              epochs=100,
                              # callbacks=callbacks,
                              verbose=1)
model_6.summary()

I've already tried using X_train, y_train, X_test, y_test directly unsuccessfully. I also used layers.Rescaling(1/255.) and a smaller learning rate but those options didn't work either. Any ideas are appreciated.

Edit 1:
I recently tried ...

from tensorflow.keras.preprocessing.image import ImageDataGenerator
train_datagen = ImageDataGenerator(rescale=1/255.)
test_datagen = ImageDataGenerator(rescale=1/255.)

train = train_datagen.flow(X_train, y_train, batch_size=BATCH_SIZE)
test = test_datagen.flow(X_test, y_test, batch_size=BATCH_SIZE)

... Instead of data API but it also didn't work.