class IoULoss(nn.Module):
    def __init__(self, weight=None, size_average=True):
        super(IoULoss, self).__init__()
        self.weights = weight

    def forward(self, inputs, targets, eps=0.001):
        inputs = nn.Softmax(dim=1)(inputs)
        targets = F.one_hot(targets.long(), 5).float()

        intersection = torch.sum((inputs.permute(0,2,3,1)*targets), (1,2))

        total = torch.sum((inputs.permute(0,2,3,1) + targets), (1,2))

        union = total - intersection

        IoU = (intersection + eps)/(union + eps)

        iouLoss = 1 - (IoU.mean(0) * self.weights).sum()/self.weights.sum()

        return iouLoss

Thank you. It worked now. But unfortunately, the model doesn’t seem to be learning from this loss.