Repository: pytorch
Issue: Forward_AD and Torchscript Functions results in Nones or wrong values.
Labels: oncall: jit,module: autograd,module: forward ad

### 🐛 Forward AD returns wrong tangent on second call  of some torchscript functions

Hi,

it seems that the combination of nested functions inside torchscript causes the forward_ad to missfunction on the second call. 
Depending on the 'bad function', the tangent returned will be either None or (for example, if stacking) zeros.
Not sure if I am holding it wrong or if this is an undocumented bug.

```python
import torch

print(torch.__version__)


@torch.jit.script
def goodfun1(x):
    return torch.sin(x)


def goodfun2(x):
    return torch.sin(x + x)


@torch.jit.script
def badfun(x):
    return torch.sin(x + x)


for desc, fun in [
    ("goodfun - no addition", goodfun1),
    ("goodfun - no torchscript", goodfun2),
    ("badfun", badfun),
]:
    print("running", desc)
    x = torch.ones(1)
    v = torch.ones(1)
    jvp1 = torch.autograd.functional.jvp(fun, x, v)[1]

    print("first run")
    with torch.autograd.forward_ad.dual_level():
        dual_input = torch.autograd.forward_ad.make_dual(x, v)
        dual_output = fun(dual_input)
        jvp2 = torch.autograd.forward_ad.unpack_dual(dual_output).tangent
    print(jvp2)
    torch.testing.assert_close(jvp1, jvp2)

    print("second run")
    with torch.autograd.forward_ad.dual_level():
        dual_input = torch.autograd.forward_ad.make_dual(x, v)
        dual_output = fun(dual_input)
        jvp2 = torch.autograd.forward_ad.unpack_dual(dual_output).tangent
    print(jvp2)
    torch.testing.assert_close(jvp1, jvp2)

    print()

```


output:

```
1.12.0.dev20220309
running goodfun - no addition
first run
tensor([0.5403])
second run
tensor([0.5403])

running goodfun - no torchscript
first run
tensor([-0.8323])
second run
tensor([-0.8323])

running badfun
first run
tensor([-0.8323])
second run
None
---------------------------------------------------------------------------

AssertionError: None mismatch: tensor([-0.8323]) is not None

```

I would expect to get for badfun the same result in both runs, like I do in te other test cases.
 
### Versions

PyTorch version: 1.12.0.dev20220309
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 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.9.5 (default, Jun  4 2021, 12:28:51)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-88-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 10.1.243
GPU models and configuration: GPU 0: NVIDIA GeForce GTX 1080 Ti
Nvidia driver version: 470.103.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.19.2
[pip3] pytorch-lightning==1.2.10
[pip3] torch==1.12.0.dev20220309
[pip3] torchaudio==0.12.0.dev20220309
[pip3] torchinfo==1.5.3
[pip3] torchmetrics==0.4.1
[pip3] torchvision==0.13.0.dev20220309
[pip3] torchviz==0.0.2
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.3.1               h2bc3f7f_2  
[conda] mkl                       2020.2                      256  
[conda] mkl-service               2.3.0            py39he8ac12f_0  
[conda] mkl_fft                   1.3.0            py39h54f3939_0  
[conda] mkl_random                1.0.2            py39h63df603_0  
[conda] numpy                     1.19.2           py39h89c1606_0  
[conda] numpy-base                1.19.2           py39h2ae0177_0  
[conda] pytorch                   1.12.0.dev20220309 py3.9_cuda11.3_cudnn8.2.0_0    pytorch-nightly
[conda] pytorch-lightning         1.2.10             pyhd8ed1ab_0    conda-forge
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                0.12.0.dev20220309      py39_cu113    pytorch-nightly
[conda] torchinfo                 1.5.3              pyhd8ed1ab_0    conda-forge
[conda] torchmetrics              0.4.1              pyhd8ed1ab_0    conda-forge
[conda] torchvision               0.13.0.dev20220309      py39_cu113    pytorch-nightly
[conda] torchviz                  0.0.2                    pypi_0    pypi

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7