Hello everyone, I am currently trying to implement custom dice loss function for my semantic segmentation model.
The problem is that I keep getting an error as below:
’RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn’
I know that this error is caused due to using torch.argmax(pred,dim=1),
as argmax function is not differentiable. However, without using argmax, I can not think of any other way as for my dice loss function I need to compare pred and target values as below:
with torch.enable_grad():
num= 2 * torch.sum(pred * target)
den = torch.sum(pred + target)
return 1 - torch.true_divide((num+ 1), (den+ 1))
Could somebody please help me with this problem?
Thank you!
My whole function looks as below:
def custom_loss(output,labels):
**    preds = torch.nn.LogSoftmax(dim=1)(output)**
**    preds = torch.argmax(preds,dim=1)**
**    print(preds)**
**    loss = dice_loss(preds,labels)**
**    return loss**
def dice_loss(pred,target):
**    with torch.enable_grad():**
**        num= 2 * torch.sum(pred * target)**
**        den = torch.sum(pred + target)**
**        return 1 - torch.true_divide((num+ 1), (den+ 1))**