Repository: pytorch
Issue: jit: Confusing behavior w/ torch.autograd.grad, iterative loop, and printing?
Labels: oncall: jit,Stale

### 🐛 Describe the bug

Context: Just trying to port some TensorFlow code to PyTorch and do some performance profiling; it does some iterative stuff with gradients of network w.r.t. inputs.

When using `torch.jit.script` with code like the following, there's some odd behavior where `num_iter >= 8` seems to change how `torch.autograd.grad` reasons about `x`... or something. For example, this code ran through jit:
```py
class FakeAlgorithm(torch.nn.Module):
    # Since we can't pass a net as argument, just wrap it into a module...
    def __init__(self, net):
        super().__init__()
        self.net = net
        self._step_size = 0.001

    def _on_iter(self, i: int):
        pass

    def _gradient(self, x):
        y = self.net(x)
        dy_dx, = torch.autograd.grad([y.sum()], [x])
        assert dy_dx is not None
        return dy_dx.detach()

    def forward(self, x, num_iter: int):
        # Bad?
        x = x.detach().requires_grad_(True)
        for i in range(num_iter):
            self._on_iter(i)
            dy_dx = self._gradient(x)
            x = x + dy_dx * self._step_size
        return x
```
Yields an error like:
```
  File "./torch_jit_issue.py", line 18, in _gradient
    def _gradient(self, x):
        y = self.net(x)
        dy_dx, = torch.autograd.grad([y.sum()], [x])
                 ~~~~~~~~~~~~~~~~~~~ <--- HERE
        assert dy_dx is not None
        return dy_dx.detach()
RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.
```

Some notes:
- It doesn't fail if not using `torch.jit.script`
- It doesn't fail if `num_iter<=7`
- It doesn't fail if (correctly?) detach inside the loop

Full script: https://github.com/EricCousineau-TRI/repro/blob/2fb2dc0db0fd6a1f848ff83f771dbc0a0b4fff4e/python/torch/torch_jit_issue.py

### Versions

```
$ python -m torch.utils.collect_env
Collecting environment information...
PyTorch version: 1.10.1+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.8.10 (default, Nov 26 2021, 20:14:08)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-5.11.0-41-generic-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration:
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090

Nvidia driver version: 470.86
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.3.1
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.17.4
[pip3] torch==1.10.1+cu111
[conda] Could not collect
```