Got "Element 0 of tensors" when using torch.where() in custom loss

George: <p>Hi all, I tried to build a custom loss function,and found that if using torch.where() in loss operation the error " Element 0 of tensors does not require grad and does not have a grad_fn" raised when I start training the model.And if I just remove torch.where() the training of model works.</p>
<pre><code class="lang-auto">class Custom_loss(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, y_hat, y):
        prob = torch.abs((y_hat - y)/y)
        t = torch.ones_like(prob)
        f = torch.zeros_like(prob)
        prob = torch.where(prob &lt; 0.1, t, f)
        return torch.sum(prob)
</code></pre>
<p>Is there a way to get around it?</p>
<p>Thanks!</p>

LeviViana: <p>I guess it might come from the fact that <code>t</code> and <code>f</code> don’t require grad.</p>
<pre><code class="lang-python">import torch
prob = torch.randn(5)
prob.requires_grad_(True)
t = torch.ones_like(prob)
f = torch.zeros_like(prob)
p = torch.where(prob &lt; 0.3, t, f)
print(p.requires_grad) # prints False
</code></pre>
<p>However if you set <code>requires_grad</code> to <code>True</code> for <code>t</code> and <code>f</code>, then:</p>
<pre><code class="lang-python">import torch
prob = torch.randn(5)
prob.requires_grad_(True)
t = torch.ones_like(prob)
f = torch.zeros_like(prob)
t.requires_grad_(True)
f.requires_grad_(True)
p = torch.where(prob &lt; 0.3, t, f)
print(p.requires_grad) # prints True
</code></pre>
<p>PS: I supposed you meant <code>f = torch.zeros_like(prob)</code> instead of <code>f = torch.ones_like(prob)</code></p>

George: <p>Thank you <a class="mention" href="/u/leviviana">@LeviViana</a>, the model can be trained after setting  <code>requires_grad</code>  to  <code>True</code>  for  <code>t</code>  and  <code>f</code>,<br>
but the loss doesn’t change,I will look into this. By the way, should I put a negative sign in front of torch.sum(prob) to make it a minimum optimization?</p>

LeviViana: <p>The problem with this snippet is that there is no gradient for the variable <code>prob</code>. Indeed, the first problem was related to the fact that your loss didn’t require gradient, and it was fixed by enabling gradients for <code>t</code> and <code>f</code>. Now the problem is that it doesn’t work right.</p>
<pre><code class="lang-python">import torch
prob = torch.rand(5)
prob.requires_grad_(True)
t = torch.ones_like(prob)
f = torch.zeros_like(prob)
t.requires_grad_(True)
f.requires_grad_(True)
p = torch.where(prob &lt; 0.3, t, f)
print(p.requires_grad) # prints True
loss = (p ** 2).sum() # Meaningless loss
loss.backward()
print(prob.grad) # No gradients
print(t.grad) # Gradients
print(f.grad) # Gradients
</code></pre>
<p>Maybe the following hack should work better for you:</p>
<pre><code class="lang-python">import torch
prob = torch.rand(5)
prob.requires_grad_(True)
max_prob = torch.max(prob)
t = torch.ones_like(prob) * max_prob
f = torch.zeros_like(prob)
p = torch.where(prob &lt; 0.3, t, f) / max_prob
print(p.requires_grad) # prints True
loss = (p ** 2).sum() # Meaningless loss
loss.backward()
print(prob.grad) # Gradients
print(t.grad) # No gradients
print(f.grad) # No gradients
</code></pre>

George: <p><a class="mention" href="/u/leviviana">@LeviViana</a>, this hack worked a few times,it’s a shame I didn’t set seed, so the success never reproduced.When I used torch.mean(prob) as substitute, the loss always decreases gradually in iterations.So I think maybe it just too hard to train the model if we turn loss into 0,1 type.</p>
<p>Thank you again  for everything you’ve done</p>

LeviViana: <p>Try analyzing the magnitude of your gradients. Maybe your loss is suffering from some numerical problems. Then, try changing the formula of your loss without changing its objective. For instance, it is well known that it’s generally better to use <code>nn.LogSofmax</code> instead of <code>nn.Softmax</code>. <a href="https://discuss.pytorch.org/t/logsoftmax-vs-softmax/21386">This thread</a> discusses this observed property.</p>