tensorflow.keras.metris で Mean IoUを算出したい

Question

tensorflowを使ってディープラーニングのセマンティックセグメンテーションのモデルをトレーニングしたのですが、モデルの評価をするためにIoU（Intersection over Union)をtensorflow.keras.metricsのライブラリを使って求めようとするとなぜかMean IoU＝１だと表示されます。トレーニングの間、Loss(binary crossentropy)は徐々に減少していたので、IoUが１だというのは間違っていると思うのですが、正しい値の算出方法がわかりません。

---

以下が実際に試してみたコードです。追加の情報が必要であればコメントをお願いします。

from tensorflow.keras.metrics import MeanIoU

#Test generator using validation data.
test_image_batch, test_mask_batch = val_img_gen.__next__()
print(test_mask_batch.shape) #(32, 256, 256, 3)

#Convert categorical to integer for visualization and IoU calculation
test_mask_batch_argmax = np.argmax(test_mask_batch, axis=3) 
test_pred_batch = (model.predict(test_image_batch)> 0.5).astype(np.uint8)
print(test_pred_batch.shape) # (32, 256, 256, 1)
test_pred_batch_argmax = np.argmax(test_pred_batch, axis=3)
print(test_mask_batch.shape) #(32, 256, 256, 3)

n_classes = 2
IOU_keras = MeanIoU(num_classes=n_classes)  
IOU_keras.update_state(test_pred_batch_argmax, test_mask_batch_argmax)
print("Mean IoU =", IOU_keras.result().numpy())

#output -- Mean IoU = 1.0