Keras report NaN for Loss during training for multiclass segmentation

Hi.

I am using this Keras tutorial to train a multiclass segmentation model. Instead of using the CIHP dataset, which has 20 classes (28280 training and 5000 validation images), I am using mut1ny’s face/head segmentation dataset, which has 14 classes (5620 Training and 1140 validation images).

Most images in mut1ny’s dataset are synthetic 3D human heads but there are real, in-the-wild human shots as well. I also have a distractor set incorporated into the dataset.

During training, this is what I get:

Epoch 1/50
 415/1404 [=======>......................] - ETA: 14:15 - loss: nan - accuracy: 0.5197

The loss is always nan and accuracy never gets better than 0.5197. The training and validation plots are parallel to the x-axis.

What am I doing wrong? How can I fix this?