Location of scripted function's definition affects result

username: _lawrence
cooked: <p>hi, I’m relatively new to torch script, and I’m seeing an effect where the <code>grad_fn</code> is not getting propagated, depending on where the scripted function is defined. Here’s the simplified example:</p>
<pre><code class="lang-auto">@torch.jit.script
def func0(x):
    return x + x + x

class SomeModule(torch.nn.Module):

    def forward(self, inp):
        # inp: has grad_fn

        @torch.jit.script
        def func1(x):
            return x + x + x

        out0 = func0(inp) # grad_fn not set
        out1 = func1(inp) # grad_fn set
</code></pre>
<p>This only represents the general idea (I’m working with a much larger system), but I’m seeing that with the same function definition, <code>grad_fn</code> gets set when the function is defined in the local scope, but not in the global scope. Even more odd, if <code>func0</code> returns only <code>x + x</code> instead of <code>x + x + x</code>, then <code>grad_fn</code> is set as expected.</p>
<p>Any ideas what could be causing these differences? Thanks! (currently using Pytorch 1.11, but same effect since at least 1.8.)</p>

username: ptrblck
cooked: <p>I’m seeing <code>grad_fn</code>s for both cases:</p>
<pre><code class="lang-python">@torch.jit.script
def func0(x):
    return x + x + x

class SomeModule(torch.nn.Module):

    def forward(self, inp):
        # inp: has grad_fn

        @torch.jit.script
        def func1(x):
            return x + x + x

        out0 = func0(inp) # grad_fn not set
        out1 = func1(inp) # grad_fn set
        
        return out0, out1
    
model = SomeModule()
inp = torch.randn(1, requires_grad=True) + 1
print(inp.grad_fn)
# &gt; &lt;AddBackward0 object at 0x7f63c5c76940&gt;
out0, out1 = model(inp)
print(out0.grad_fn)
# &gt; &lt;CppFunction at 0x7f63c5c76c10&gt;
print(out1.grad_fn)
# &gt; &lt;AddBackward0 at 0x7f63c5c76dc0&gt;
</code></pre>
<p>in <code>1.11.0.dev20211101</code>. Could you update to the latest nightly binary and rerun the code?</p>

username: anantguptadbl
cooked: <p>I tried it in a slightly different manner and it works</p>
<pre><code class="lang-auto">import os
import numpy as np
import time
import torch
import torchvision
from torch import nn
from torch.autograd import Variable

def func0(x, y):
    return x + y

class SomeModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.y = nn.Parameter(torch.randn([1]))

    def forward(self, inp):
        def func1(x):
            return x + self.y

        out0 = func0(inp, self.y)
        out1 = func1(inp)
        return out0, out1
        
model = SomeModule()
inp = torch.randn(1) + 1
print(inp.grad_fn)
out0, out1 = model(inp)
print(out0.grad_fn)
print(out1.grad_fn)
</code></pre>
<pre><code class="lang-auto">torch.__version__
'1.8.1'
</code></pre>

username: _lawrence
cooked: <p>Thanks for your responses. I probably wasn’t clear, that this isn’t the actual program I’m running (this simplified version works for me also), but rather I was using it to illustrate that there’s a difference between whether the scripted function is defined in the global scope versus local scope. Adding on to this, it doesn’t work either if the function is defined inside <code>def __init__()</code> and then set with <code>self.func = func</code>, and then later called in <code>forward</code> … the only case that works is when the scripted function is defined within <code>forward</code>.</p>
<p>One other detail I noticed is that the global case works if it only performs a single operation (e.g., <code>x + x</code>), rather than 2+ operations (e.g., <code>x + x + x</code>), which leads me to believe this might be a bug in the fuser.</p>
<p>Also, I’m not sure, but this topic seems possibly related: <a href="https://discuss.pytorch.org/t/second-forward-call-of-torchscripted-module-breaks-on-cuda/124291" class="inline-onebox">Second forward call of torchscripted module breaks on cuda</a></p>
<p>thanks again!</p>