Repository: pytorch
Issue: `jit(Function)` results in double execution
Labels: oncall: jit

### 🐛 Describe the bug

`jit`-compiling a function that calls a custom function (`torch.autograd.Function`) results in the following behavior:

* The contents of the custom function are added to TorchScript but the output of the jitted code is discarded
* The Python version of the custom function will always be called before the TorchScript version and its output will be used in later code.

So effectively, the custom function is executed twice!

To reproduce:

```python
import torch


class MyFunction(torch.autograd.Function):

    @staticmethod
    def forward(ctx, *args, **kwargs):
        print("Python: ", end=" ")
        x = print_executing_forward(args[0])
        return x + 1

    @staticmethod
    def backward(ctx, *grad_args):
        pass


my_function = MyFunction().apply


@torch.jit.script_if_tracing
def print_executing_forward(x: torch.Tensor):
    print("Executing forward()")
    return x


def script(x):
    print("Tracing script...")
    return my_function(x)


print("Tracing script_jit ...")
script_jit = torch.jit.trace(script, torch.ones(()))

for i in range(10):
    print(f"Running script_jit for the {i + 1}th time ...")
    script_jit(torch.ones(()))
    print()
```

This results in the following output (cropped):
```
[...]

Running script_jit for the 9th time ...
Python:  Executing forward()
Executing forward()

Running script_jit for the 10th time ...
Python:  Executing forward()
Executing forward()
```

### Versions

Collecting environment information...
PyTorch version: 1.11.0+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
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.9.7 (default, Sep 16 2021, 13:09:58)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.13.0-35-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3090
Nvidia driver version: 470.103.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.2.1
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.5
[pip3] torch==1.11.0+cu113
[pip3] torchaudio==0.10.2+cu113
[pip3] torchvision==0.11.3+cu113
[conda] blas                      1.0                         mkl
[conda] cpuonly                   1.0                           0    pytorch
[conda] cudatoolkit               11.3.1               h2bc3f7f_2
[conda] mkl                       2022.0.1           h06a4308_117
[conda] numpy                     1.21.5                   pypi_0    pypi
[conda] pytorch-mutex             1.0                         cpu    pytorch
[conda] torch                     1.11.0+cu113             pypi_0    pypi
[conda] torchaudio                0.10.2+cu113             pypi_0    pypi
[conda] torchvision               0.11.3+cu113             pypi_0    pypi