IOU Loss function implementation in Pytorch

Hello All,

I am trying to implement IOU as loss function for my semantic segmentation problem which has multiple classed. I am looking for pytorch implementation and found the post Understanding different Metrics implementations (IoU) but it does not support multiple class. Can anyone point me towards IOU Pytorch implementation implementation as loss function? Thanks