Modified by an inplace operation only if JIT compiled

avitase: <p>Hi:)</p>
<p>I have written an <code>nn.Module</code> that works fine if used as a plain Python object but as soon as I <code>jit.script</code> / compile it, it throws at runtime:</p>
<pre><code class="lang-auto">RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.DoubleTensor [10, 8]], which is output 0 of SelectBackward, is at version 3; expected version 2 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!
</code></pre>
<p>Why does this only occur after compilation? Naively, I would understand if such an error occur if <strong>not</strong> compiled and then gets <em>optimized</em> away by compiling. But here the compiler seemingly introduces this error. Could this be a bug in the compiler?</p>
<p>Thanks for your help!</p>
<p><strong>EDIT</strong><br>
I managed to bisect my code and reduce it to this snippet:</p>
<pre><code class="lang-auto">debug = False

model = MyFancyModel()
if not debug:
    model = torch.jit.script(model)
loss = torch.sum((model(x) - y)**2)
print('#1', loss)
loss.backward()  # works
print('#1 done.')

model = MyFancyModel()
if not debug:
    model = torch.jit.script(model)
loss = torch.sum((model(x) - y)**2)
print('#2', loss)
loss.backward()  # fails
print('#2 done.')
</code></pre>
<p>which works fine if <code>debug</code> is set to <code>True</code> but fails on the second call of <code>loss.backward()</code>. I am very confused. <code>MyFancyModel()</code> shouldn’t carry any static state and I don’t understand why (a) the second invocation of <code>.backward()</code> works w/o compiling but fails if the module is compiled, and (b) why does the first invocation works even with compiling!? Do you have any suggestions how to attack this problem?</p>

jmandivarapu1: <p>I faced the same error recently in different scenario(not in jit) and I found the solution like below</p>
<pre><code class="lang-auto">      outer_loss = torch.tensor(0., device=args.device)
     avg_loss=F.cross_entropy(test_logit, test_target)
     outer_loss += avg_loss

     with torch.no_grad(): #if i don't change the loss using this it's throwing error as the model states for the claculated losses are different
              outer_loss.set_(torch.Tensor([InitialLoss-F.cross_entropy(test_logit, test_target)]).to(device)) 
</code></pre>
<p>This link also helped in understanding the problem</p>

avitase: <p>Thanks a lot for your help. But I don’t see the connection to my problem. Do you suggest to add <code>with torch.no_grad()</code> to my code? But where?</p>

jmandivarapu1: <aside class="quote no-group" data-username="avitase" data-post="1" data-topic="121635">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/94ad74/48.png" class="avatar"> avitase:</div>
<blockquote>
<pre><code class="lang-auto">debug = False

model = MyFancyModel()
if not debug:
    model = torch.jit.script(model)
loss = torch.sum((model(x) - y)**2)
print('#1', loss)
loss.backward()  # works
print('#1 done.')

model = MyFancyModel()
if not debug:
    model = torch.jit.script(model)
loss = torch.sum((model(x) - y)**2)
print('#2', loss)
loss.backward()  # fails
print('#2 done.')
</code></pre>
</blockquote>
</aside>
<p>I referred about the error can you paste you MyFancyModel() for debugging</p>

tom: <blockquote>
<p>why (a) the second invocation of <code>.backward()</code> works w/o compiling but fails if the module is compiled, and<br>
(b) why does the first invocation works even with compiling!? Do you have any suggestions how to attack this problem?</p>
</blockquote>
<p>Very likely, one of the inputs is modified in place but the JIT fuser combines operations in a way that it needs the input for the backward. This typically kicks in with the second run on (the first being an information gathering (profiling the shapes) run).</p>
<p>You could try if anomaly detection helps you hunt down the inplace modification and remove that.</p>

avitase: <p>The model is quite involved. I will try to reduce it s.t. it becomes more readable. Thanks again:)</p>

avitase: <p>The error message<br>
<code>RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.DoubleTensor [10, 8]], which is output 0 of SelectBackward, is at version 3; expected version 2 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!</code><br>
is already the output of the anomaly detection.<br>
Are you referring to the <code>x</code> and <code>y</code> in the example? If my module works correctly these variables should be read-only, but I will double check. All other variables should be renewed by the second compilation step.</p>
<p><strong>EDIT</strong><br>
I am now quite sure that <code>x</code> and <code>y</code> do not change. I copy both via <code>tensor.clone()</code> and compare them after calling my module. All other variables are renewed by calling <code>model = MyFancyModel()</code> the second time…</p>

tom: <p>Do you use select and indexing somewhere and then change the tensor you indexed?<br>
As those will share memory (they are 'view’s into the same memory in autograd lingo) they share the same version counter and modifying the “parent” tensor will invalidate using the view in the backward (as it doesn’t track which part was viewed and which part was modified, that is the case regardless of the actual values have changed).</p>

avitase: <p>Yes, I do use indexing. I managed to strip down my module code, unrolled a few loops and ended up with this snippet:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn


class Core(nn.Module):
    def __init__(self) -&gt; None:
        super().__init__()
        self.f = nn.Linear(4, 2)

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -&gt; torch.Tensor:
        x = torch.cat((x1, x2), 1)
        return self.f(x)


class MyFancyModel(nn.Module):
    def __init__(self) -&gt; None:
        super().__init__()
        self.core = Core()

    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
        assert x.shape == (1, 2, 2)

        y = torch.empty(1, 4, 2)
        y[:, :2] = x

        x1, x2 = x[:, 0], x[:, 1]
        y1, y2 = y[:, 0], y[:, 1]
        y[:, 2] = x2 + (x2 - x1) * self.core(y1, y2)

        x1, x2 = x[:, 1], y[:, 2]
        y1, y2 = y[:, 1], y[:, 2]
        y[:, 3] = x2 + (x2 - x1) * self.core(y1, y2)

        return y
</code></pre>
<p>In case my indexing is broken, can you give me a hint how to fix it?</p>

tom: <p>You could try sticking your y bits in a list and then use torch.cat.<br>
So something like this (but I might have gotten something wrong).<br>
<code>x[:, None]</code> is the same as <code>x.unsqueeze(1)</code>, that is my taste preference here, you might have different ones.</p>
<pre><code class="lang-python">class MyFancyModel(nn.Module):
    def __init__(self) -&gt; None:
        super().__init__()
        self.core = Core()

    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
        assert x.shape == (1, 2, 2)

        y = torch.empty(1, 4, 2)
        y_list = [x]

        x1, x2 = x[:, 0], x[:, 1]
        y1, y2 = x[:, 0], x[:, 1]
        y_new = x2 + (x2 - x1) * self.core(y1, y2)
        print(y_new.shape)
        y_list.append(y_new[:, None])

        x1, x2 = x[:, 1], y_new # ??
        y1, y2 = y2, y_new
        y_new = x2 + (x2 - x1) * self.core(y1, y2)
        y_list.append(y_new[:,None])
        y = torch.cat(y_list, 1)

        return y
</code></pre>
<p>Best regards</p>
<p>Thomas</p>

avitase: <p>Thanks for your help. Much appreciated! So you are saying assigning values to a tensor as in <code>y[:, 2] = ...</code> does not work with autograd, even if this value is never overwritten afterwards?</p>

tom: <p>Yeah. I have this superfancy “everything about autograd course” that covers this error in full detail. Writing and recording the course and all was great fun, but I am slow to set up the delivery…</p>

avitase: <p>thanks for your help:)</p>