Can't use Keras MeanIoU to train semantic segmentation model

I'm working on a binary semantic segmentation problem. I built an UNet model with MobileNetV2 backbone. Here is my model code:

def upsample(filters, size, apply_dropout=False):
    initializer = tf.random_normal_initializer(0., 0.02)
    layer = Sequential()
    layer.add(layers.Conv2DTranspose(filters, size, strides=2, padding='same', kernel_initializer=initializer,
                                     use_bias=False))
    layer.add(layers.BatchNormalization())
    if apply_dropout:
        layer.add(layers.Dropout(0.5))
    layer.add(layers.ReLU())
    return layer


def UNet(image_size, num_classes):
    inputs = Input(shape=image_size + (3,))

    base_model = applications.MobileNetV2(input_shape=image_size + (3,), include_top=False)
    layer_names = [
        'block_1_expand_relu',
        'block_3_expand_relu',
        'block_6_expand_relu',
        'block_13_expand_relu',
        'block_16_project',
    ]
    base_model_outputs = [base_model.get_layer(name).output for name in layer_names]
    down_stack = Model(inputs=base_model.input, outputs=base_model_outputs)
    down_stack.trainable = False

    up_stack = [
        upsample(512, 3),
        upsample(256, 3),
        upsample(128, 3),
        upsample(64, 3)
    ]

    skips = down_stack(inputs)
    x = skips[-1]
    skips = reversed(skips[:-1])

    for up, skip in zip(up_stack, skips):
        x = up(x)
        x = layers.Concatenate()([x, skip])

    outputs = layers.Conv2DTranspose(filters=num_classes, kernel_size=3, strides=2, padding='same')(x)

    return Model(inputs, outputs)

To load the images and masks for training, I built an image loader inherits from keras.Sequnce.

class ImageLoader(utils.Sequence):

    def __init__(self, batch_size, img_size, img_paths, mask_paths):
        self.batch_size = batch_size
        self.img_size = img_size
        self.img_paths = img_paths
        self.mask_paths = mask_paths

    def __len__(self):
        return len(self.mask_paths) // self.batch_size

    def __getitem__(self, idx):
        i = idx * self.batch_size
        batch_img_paths = self.img_paths[i:i + self.batch_size]
        batch_mask_paths = self.mask_paths[i:i + self.batch_size]

        x = np.zeros((self.batch_size,) + self.img_size + (3,), dtype='float32')
        for j, path in enumerate(batch_img_paths):
            img = utils.load_img(path, target_size=self.img_size)
            img = utils.img_to_array(img)
            x[j] = img

        y = np.zeros((self.batch_size,) + self.img_size + (1,), dtype='uint8')
        for j, path in enumerate(batch_mask_paths):
            img = utils.load_img(path, target_size=self.img_size, color_mode='grayscale')
            img = utils.img_to_array(img)
            # [0, 255] -> [0, 1]
            img //= 255
            y[j] = img

        return x, y

In my segmentation problem, all the labels are in the range [0, 1]. However, when I try to compile and then fit the model using Adam optimizer, Sparse categorical cross entropy loss and metric tf.keras.metrics.MeanIoU, I encountered with the following problem:

Node: 'confusion_matrix/assert_non_negative_1/assert_less_equal/Assert/AssertGuard/Assert'
2 root error(s) found.
  (0) INVALID_ARGUMENT:  assertion failed: [`predictions` contains negative values.  ] [Condition x >= 0 did not hold element-wise:] [x (confusion_matrix/Cast:0) = ] [-1 -1 -1...]
     [[{{node confusion_matrix/assert_non_negative_1/assert_less_equal/Assert/AssertGuard/Assert}}]]
     [[confusion_matrix/assert_less_1/Assert/AssertGuard/pivot_f/_31/_67]]
  (1) INVALID_ARGUMENT:  assertion failed: [`predictions` contains negative values.  ] [Condition x >= 0 did not hold element-wise:] [x (confusion_matrix/Cast:0) = ] [-1 -1 -1...]
     [[{{node confusion_matrix/assert_non_negative_1/assert_less_equal/Assert/AssertGuard/Assert}}]]

At first, I used accuracy as a metrics for training and I didn't encounter this problem, however when I changed to MeanIoU, this problem happened. Does anyone know how to fix this problem? Thank you very much!

UPDATE: I've searched on StackOverflow and found this question about a similar error, however the fix mentioned in that link (reduce learning rate) doesn't work in my case.

How many English words
do you know?
Test your English vocabulary size, and measure
how many words do you know
Online Test
Powered by Examplum