Jit.script and torch.no_grad(): Different behavior with and without JIT?

username: tjoseph
cooked: <p>Hi,<br>
can someone explain if this behavior is intended and if so, why and what the solution would be to get the same behavior with scripting as without:</p>
<pre><code class="lang-auto">import torch


@torch.jit.script
def scripted_bar(x):
    with torch.no_grad():
        y = x * 2

    z = 2 * x + y
    print("scripted_bar", x.requires_grad, y.requires_grad, z.requires_grad)

    return z


def bar(x):
    with torch.no_grad():
        y = x * 2

    z = 2 * x + y
    print("bar", x.requires_grad, y.requires_grad, z.requires_grad)

    return z


a = torch.rand(3, requires_grad=True)
b = torch.rand(3, requires_grad=True)


print("scripted_bar(a)", scripted_bar, scripted_bar(a).requires_grad)
print("bar(a)", bar, bar(a).requires_grad)
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">scripted_bar True False False
scripted_bar(a) &lt;torch.jit.ScriptFunction object at 0x7fa72315d450&gt; False
bar True False True
bar(a) &lt;function bar at 0x7fa797fae3a0&gt; True
</code></pre>
<p>The <code>z</code> variable does have no gradient tracking with <code>jit.script</code>.</p>

username: ptrblck
cooked: <p>This behavior seems indeed unexpected. Could you create a <a href="https://github.com/pytorch/pytorch/issues">GitHub issue</a> with this minimal code snippet, please?</p>

username: tjoseph
cooked: <p>Thank you for answering <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"> I will create an issue</p>
<aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/75788">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/75788" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/75788" target="_blank" rel="noopener nofollow ugc">`torch.jit.script` Script functions do return `requires_grad = False` if `torch.no_grad()` has been used</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2022-04-14" data-time="09:50:58" data-timezone="UTC">09:50AM - 14 Apr 22 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/mctigger" target="_blank" rel="noopener nofollow ugc">
          <img alt="mctigger" src="https://avatars.githubusercontent.com/u/4086336?v=4" class="onebox-avatar-inline" width="20" height="20">
          mctigger
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          oncall: jit
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">### 🐛 Describe the bug

Hello,
I think I found a bug with `torch.jit.script` an<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">d `torch.no_grad()`. Here is a code sample to reproduce behavior:

```
import torch


def bar(x):
    with torch.no_grad():
        y = x * 2

    z = 2 * x + y
    print("requires_grad: ", x.requires_grad, y.requires_grad, z.requires_grad)

    return z


scripted_bar = torch.jit.script(bar)

a = torch.rand(3, requires_grad=True)
b = torch.rand(3, requires_grad=True)

print("scripted_bar(a) : ", scripted_bar, scripted_bar(a).requires_grad)
print("bar(b) : ", bar, bar(b).requires_grad)

```

I would expect the scripted and non-scripted function to both return tensors with `requires_grad = True`.


### Versions

Collecting environment information...
PyTorch version: 1.11.0+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.8.13 (default, Mar 28 2022, 11:38:47)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.13.0-39-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: Quadro T1000
Nvidia driver version: 510.54
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.3
[pip3] torch==1.11.0
[pip3] torchvision==0.12.0
[conda] mypy-extensions           0.4.3                    pypi_0    pypi
[conda] numpy                     1.22.3                   pypi_0    pypi
[conda] torch                     1.11.0                   pypi_0    pypi
[conda] torchvision               0.12.0                   pypi_0    pypi</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>