Detach but still make autograd trace the gradients

Abdelrahman_Akram: <p>In the middle of my model, there is some for loops that I <code>must</code> do, there is no replace for it, so to make it faster I will use numba.njit, but it works for only numpy arrays and also if I detach the tensor the gradients won’t be recorded. So is there a solution for that, and also I heard of torch.jit.trace, will it do the same as numba.njit and record the gradients.</p>

albanD: <p>Hi,</p>
<p>If you have to break the autograd, you can use a custom Function <a href="https://pytorch.org/docs/stable/notes/extending.html" class="inline-onebox">Extending PyTorch — PyTorch 1.8.1 documentation</a>.<br>
You will have to specify what the backward pass of this op is though.</p>

Abdelrahman_Akram: <p>Oh I get it, another thing please, in this code</p>
<pre><code class="lang-auto">import torch
depth  = torch.randn((22000,3), requires_grad= True)
depth[depth&lt;0].zero_()

rot_sin = torch.sin(torch.FloatTensor([0.4]))
rot_cos = torch.cos(torch.FloatTensor([0.4]))
rot_mat_T = torch.FloatTensor(
            [[rot_cos, 0, -rot_sin], [0, 1, 0], [rot_sin, 0, rot_cos]],
        )
rot_mat_T.requires_grad = True
depth = depth @ rot_mat_T

linear = torch.nn.Linear(3,4)
out = linear(depth)
out.sum().backward(retain_graph =True)
print(depth.grad)
</code></pre>
<p>It gives me this warning</p>
<pre><code class="lang-auto">UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the gradient for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations.
  print(depth.grad)
</code></pre>
<p>put if I replaced this line <code>depth = depth @ rot_mat_T</code> with <code>new_depth = depth @ rot_mat_T</code><br>
and <code>out = linear(depth)</code> to <code>out = linear(new_depth)</code> it works and print the gradient, why is that?</p>

albanD: <aside class="quote no-group" data-username="Abdelrahman_Akram" data-post="3" data-topic="122102">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/abdelrahman_akram/48/38628_2.png" class="avatar"> Abdelrahman_Akram:</div>
<blockquote>
<p><code>depth = depth @ rot_mat_T</code></p>
</blockquote>
</aside>
<p>When you do that, you make the “depth” python variable point to the result of <code>depth @ rot_mat_T</code> instead of the Tensor it was pointing out before (the result of <code>torch.randn((22000,3), requires_grad= True)</code>).</p>
<p>But the <code>.grad</code> field is only populated for leaf Tensors that require gradients.<br>
And you can check that the “depth” just afterr creating it at the top .is_leaf==True.<br>
But after <code>depth = depth @ rot_mat_T</code> the new Tensor “depth” points to has .is_leaf==False.</p>
<p>So the .grad field won’t be populated, hence the warning.</p>

Abdelrahman_Akram: <p>So if I want the grads flow to the first depth variable, I have to make the new_depth step, right ?<br>
and I f I did so won’t it take more memory !?</p>

albanD: <p>Not really, because the computational graph keeps a reference to the original “depth” Tensor and so that Tensor stays alive anyways. So you won’t see any memory difference.</p>