I’m training a network using a custom loss function. And it gives a RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn. It seems that the indices from max is not differentiable which i have to use for loss function.
Does anyone have idea to fix it?
This is my code:
class MyLoss(nn. Module):
    def __init__(self):
        super(MyLoss, self).__init__()
        
    def forwad(self, d, m_label, n_label):
        # d size : [N, C, H, W], from softmax output
        # m_label size : [N, 1, H, W]
        # n_label size : [N, 1, H, W]
        
        # get the predicted class
        pred = d.max(1)[1]
        pred = pred.unsqueeze(1)
        n_pred = 2.7*pred + m_label
        
        # calc image grad along x, y
        m_dx = m_label[:, :, 1:, :] - m_label[:, :, :-1, :]
        m_dy = m_label[:, :, :, 1:] - m_label[:, :, :, :-1]
        n_dx = n_pred[:, :, 1:, :] - n_pred[:, :, :-1, :]
        n_dy = n_pred[:, :, :, 1:] - n_pred[:, :, :, :-1]
        
        # loss1: minimize the difference of grad of m_label and n_pred
        loss1 = torch.mean(((n_dx - m_dx).abs().sum() 
                              + (n_dy - m_dy).abs().sum()))
        
        # loss2: minimize the difference of n_pred and n_label
        loss2 = nn.functional.l1_loss(n_pred, n_label)
        return loss1 + loss2