hi, I’m relatively new to torch script, and I’m seeing an effect where the grad_fn is not getting propagated, depending on where the scripted function is defined. Here’s the simplified example:
@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

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