The loss works as fine as cross entropy when l_one = 0.
However, it does not decrease when I pass a zero matrix. It should behave the same as above because the custom term is multiplied by zero. This is validated by the fact that my assertion statements pass. @tom does autograd.grad influence how the optimizer updates the model parameters, or is this a potential bug for autograd?

def explanation_loss(logits, x, target, \
        mask_matrix=torch.zeros((len(dataset), 3, 32, 32)), \
        l_one=0.001):

    assert x.requires_grad == True
    ce = F.cross_entropy(logits, target) 
    if l_one == 0:
        return ce

    sum_log_hat_y = torch.sum(F.log_softmax(logits, dim=1)).mean()
    # same dims as x
    input_grad, = torch.autograd.grad(sum_log_hat_y, x, retain_graph=True, create_graph=True)
    
    dot_prod = torch.dot(torch.flatten(mask_matrix), torch.flatten(input_grad))
    assert dot_prod == 0, dot_prod.item()
    
    right_reasons = l_one * ((dot_prod) ** 2)
    assert right_reasons == 0, right_reasons.item()  # These assertions pass
    
    return ce + right_reasons