`copy_` operations get repeated in autograd computation graph

username: fxmarty
cooked: <p>Hello,</p>
<p>In part of my model, a tensor <em>that is not parameters</em> gets modified during the <code>forward()</code> call and I want to remember it. As I use DataParallel, this tensor is registered as a buffer so that it will be saved from the first replica, as explained here in the second warning: <a href="https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html" class="inline-onebox" rel="noopener nofollow ugc">DataParallel — PyTorch 2.1 documentation</a></p>
<p>The modification to this tensor <em>needs to be in-place</em>. Hence, I want to use <code>copy_</code>.</p>
<p>This is from there that I found a weird behavior for <code>copy_</code>, so I put the code below just as a sample displaying it (happening without DataParallel).</p>
<p>Here is my sample code:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.optim as optim

from torchviz import make_dot


class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        
        self.lin = nn.Linear(4, 8, bias=False)
        
        self.lin_bis = nn.Linear(8, 4, bias=False)
        
        self.register_buffer('tensor_bidon', torch.ones(self.lin.weight.shape))
        
    def forward(self, x):
        self.tensor_bidon.copy_(self.lin.weight)
         
        x = self.lin(x)
        x = self.lin_bis(x)
                
        x = torch.matmul(x, self.tensor_bidon.t())
        
        return x


mymod = MyModel()

loss_mine = torch.nn.MSELoss()

optimizer_mine = optim.Adam(mymod.parameters())

mymod.train()

for i in range(5):
    print(f'-------- {i} ---------')
    optimizer_mine.zero_grad()
    inp = torch.randn(10, 4)
    target = torch.randn(10, 8)

    outputs_full = mymod(inp)

    loss = loss_mine(outputs_full, target)
    
    dot = make_dot(outputs_full, params={**{'inputs': inp}, **dict(mymod.named_parameters())})
    dot.render(f"MYDIAG_SINGLE{i}", format='png')
    
    print(mymod.lin.weight)
    loss.backward()

    optimizer_mine.step()
</code></pre>
<p>Here is the first autograd graph forward (as saved with <code>make_dot</code> in my code above):</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/f/7/f72d1ba1f8abbdd3f49fd15cdde7e6c775226197.png" alt="MYDIAG_SINGLE0" data-base62-sha1="zgCzROPVM5AfbchtwVPyOm9MPcj" width="463" height="381"></p>
<p>As you can see (in the posts below), the CopyBackwards operation gets duplicated in the following iterations.</p>
<p>Do you have any hint as to why this happens?</p>
<p>Many thanks!</p>

username: fxmarty
cooked: <p><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/a/7afdf2db2bd19c9c7d385cf14c386bcdad2a100b.png" alt="MYDIAG_SINGLE1" data-base62-sha1="hy2kHrAuq9Aj615EiwfhZ9h9ZvJ" width="515" height="381"></p>

username: fxmarty
cooked: <p><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/0/70e5f0fdb752c21760b3094902a5a04382b36e4d.png" alt="MYDIAG_SINGLE2" data-base62-sha1="g6K7FUxMGo0XBVg4CPbGERaivut" width="575" height="459"></p>
<p>Ideally, we would not want to have those duplicates CopyBackwards.</p>

username: albanD
cooked: <p>Hi,</p>
<p>The problem is that because you copy something the requires grad into this, that Tensor starts requirering gradients, and so you start tracking all operations that happen to it.</p>
<p>In general, I would recommend to just have it being a temporary variable in the forward so that you don’t need to worry about that.<br>
If for some reason, you want to keep it on <code>self.</code>, you can detach it from the previous iterations by doing <code>self.tensor_bidon = self.tensor_bidon.detach()</code>.</p>

username: fxmarty
cooked: <p>Thanks a lot <a class="mention" href="/u/alband">@albanD</a> . That is what I ended up doing indeed, avoid the <code>copy_</code> for tensors in the gradient computation graph, and using it only in an additional buffer to save the tensor I need, where we use <code>copy_(whatever_tensor.detach())</code> (because of DataParallel requirements).</p>