Can't use Keras MeanIoU to train semantic segmentation model

Question

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.