Intersection over union (IOU) metric for multi-class semantic segmentation task

Question

I have a semantic segmentation task to predict 5 channel mask using UNET for example mask shape is (224,244,5).

I\'m using this function for IOU :

    def