Hi,
I am trying to build a U-Net Multi-Class Segmentation model for the brain tumor dataset. I implemented the dice loss using nn.module and some guidance from other implementations on the internet. But during my training, my loss is fluctuating and not converging. If I train my model using CrossEntropyLoss it is converging well. When I was debugging with the required_gradient it seems to be False for the output from the loss. I am unable to find out the issue here.
Thanks in Advance.
The loss implementation:
class DiceLoss(nn.Module):
    def __init__(self, weight=None, size_average=True):
        super(DiceLoss, self).__init__()
        self.weights = weight

def forward(self, inputs, targets, eps=0.001):
        
        inputs = torch.argmax(F.log_softmax(inputs, dim=1), dim=1)
        inputs = F.one_hot(inputs, 5).float()
        targets = F.one_hot(targets, 5).float()
        intersection = (inputs.view(4,-1,5) * targets.view(4,-1,5)).sum(1)
        total = (inputs.view(4,-1,5) + targets.view(4,-1,5)).sum(1)
        dice = (2 * intersection + eps)/(total + eps)
        diceLoss = 1 - (dice.mean(0) * self.weights).sum()/self.weights.sum()

        return diceLoss