Hi everyone,I came across some problems about gradient update when training my network.i’m using u-net as my model,and i write a custom loss function as Dice Loss Function,
class DiceCoffLoss(nn.Module):

    def __init__(self, label_of_interest=1):
        super(DiceCoffLoss, self).__init__()
        self.labelInterest = label_of_interest

    def forward(self,prediction, segmentation ):
        
        """ inputs are 2d arrays """
        segmentation=torch.autograd.Variable(segmentation) 
        prediction=torch.autograd.Variable(prediction) 
       
        if prediction.shape != segmentation.shape:
            raise ValueError("Shape mismatch between given arrays. prediction %s vs segmentation %s" \
                         % (str(prediction.shape), str(segmentation.shape)))

        n_organ_seg =  (segmentation==self.labelInterest).sum()
        n_organ_pred= (prediction == self.labelInterest).sum()
        denominator = n_organ_pred + n_organ_seg
        if denominator == 0:
            return torch.tensor(1.0,requires_grad=True) #mask or predected is empty of interested label
        
        iflat = prediction.contiguous().view(-1)
        tflat = segmentation.contiguous().view(-1)
        organ_intersection=(iflat==self.labelInterest)*(tflat==self.labelInterest) #Subscription operator
        n_organ_intersection = organ_intersection.sum()
        dice = ((2.0*n_organ_intersection / denominator) ) 
        return torch.tensor(1-dice ,requires_grad=True)

i used my model with other Loss functions like nn.functional.cross_entropy everything was OK ,but when i’m using my own loss function loss.backward can not update model and model can’t train anymore.
I am still working on that and failed to figure out why the gradient didn’t update as usual. Any inspiration would be sincerely appreciated!!