Repository: pytorch
Issue: requires_grad does not get propagated properly when using the JIT compiler
Labels: module: autograd,triaged

## 🐛 Bug

I encountered really strange behavior when using the JIT compiler to speed up my code. In particular, I noticed that when using the JIT compiler, the backward computation graph was missing some branches compared to the non-JIT compiled version. In the code below I isolated a minimal example in which the problem appears to occur: 

## To Reproduce

Steps to reproduce the behavior:

```python
from typing import Tuple, List

import torch
from torch.autograd import Variable
from torch.jit import ScriptModule


class Mod(ScriptModule):
    def __init__(self):
        super(Mod, self).__init__()
        self.var = Variable(torch.randn(2), requires_grad=True)

    @torch.jit.script_method
    def forward(self, input: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        var_lst = [self.var]
        var = torch.cat(var_lst)
        output = var + input
        print("forward: var_lst[0].requires_grad = {}".format(var_lst[0].requires_grad))
        return output, var_lst


mod = Mod()

for i in range(2):
    print("\nIteration {}".format(i))
    input = torch.randn((2,))
    output, var_lst = mod.forward(input)
    print("main: var_lst[0].requires_grad = {}".format(var_lst[0].requires_grad))
    loss = torch.cat(var_lst).sum()
    loss.backward()
```

In this example there is one variable `var` encapsulated in the `ScriptModule` `Mod` that has `requires_grad=True` set. The variable is used to compute two outputs, `output` and `var_lst`, which is just the variable wrapped in a python list. Hence, both outputs depend on the variable and should hence have `requires_grad` set to true. However, if I run the above code, I get the following output:
```
Iteration 0
forward: var_lst[0].requires_grad = True
main: var_lst[0].requires_grad = True

Iteration 1
forward: var_lst[0].requires_grad = True
main: var_lst[0].requires_grad = False
```
followed by an error, because `loss.backward()` cannot be computed when `loss.requires_grad` is false.

Interestingly, it seems that changing anything in this example leads to the error disappearing. For example removing the computation of `output` or returning `output, [self.var]` instead of `output, var_lst` causes `var_lst[0].requires_grad` to be true again.

## Expected behavior

If the code is run with `PYTORCH_JIT=0`, everything works as expected and no error occurs:
```
Iteration 0
forward: var_lst[0].requires_grad = True
main: var_lst[0].requires_grad = True

Iteration 1
forward: var_lst[0].requires_grad = True
main: var_lst[0].requires_grad = True
```

## Environment

Collecting environment information...
PyTorch version: 1.8.1
Is debug build: False
CUDA used to build PyTorch: 11.2
ROCM used to build PyTorch: N/A

OS: Arch Linux (x86_64)
GCC version: (GCC) 10.2.0
Clang version: 11.1.0
CMake version: version 3.20.0

Python version: 3.9 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 10.1.243
GPU models and configuration: GPU 0: GeForce GTX 1070
Nvidia driver version: 460.67
cuDNN version: Probably one of the following:
/usr/lib/libcudnn.so.8.1.0
/usr/lib/libcudnn_adv_infer.so.8.1.0
/usr/lib/libcudnn_adv_train.so.8.1.0
/usr/lib/libcudnn_cnn_infer.so.8.1.0
/usr/lib/libcudnn_cnn_train.so.8.1.0
/usr/lib/libcudnn_ops_infer.so.8.1.0
/usr/lib/libcudnn_ops_train.so.8.1.0
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.20.1
[pip3] torch==1.8.1
[pip3] torch-scatter==2.0.6
[conda] Could not collect

Thanks a lot in advance,
Tim



cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7 @gmagogsfm