Repository: pytorch
Issue: JIT model with mean will crash when computing the gradients on cuda
Labels: triage review,oncall: jit,module: crash

### 🐛 Describe the bug

The JIT model with mean will crash when computing the gradients on cuda.

It seems that it frees an invalid pointer

```py
import torch

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.dim = -1
        self.keepdim = False

    def forward(self, input):
        dim = self.dim
        keepdim = self.keepdim
        input = torch.sub(input, torch.tensor(9, dtype=torch.float32, device='cuda'))
        fn_res = torch.mean(input, dim, keepdim=keepdim, )
        fn_res = torch.sub(fn_res, torch.tensor(-3, dtype=torch.float32, device='cuda'))
        return fn_res

fn = M().to('cuda')
torch.random.manual_seed(56288)
inp = torch.empty([1, 64], dtype=torch.float32, memory_format=torch.contiguous_format)
inp.uniform_(-128, 63)

def clone(tensor):
    return tensor.clone().to('cuda').requires_grad_()

jit_fn = torch.jit.trace(fn, clone(inp))

from torch.autograd.functional import jacobian
jit_fn(clone(inp))

jacobian(jit_fn, (clone(inp), ))

jacobian(jit_fn, (clone(inp), ), vectorize=True, strategy='forward-mode')
```

```
free(): invalid pointer
[1]    3996840 IOT instruction (core dumped)  python mean.py
```

### Versions

```
PyTorch version: 1.13.0.dev20221001
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.1 LTS (x86_64)
GCC version: (Ubuntu 11.2.0-19ubuntu1) 11.2.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.9.13 (main, Aug 25 2022, 23:26:10)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.19.5-051905-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 11.7.99
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3060
Nvidia driver version: 515.65.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.23.3
[pip3] torch==1.13.0.dev20221001
[pip3] torchaudio==0.13.0.dev20221001
[pip3] torchvision==0.14.0.dev20221001
[conda] blas                      1.0                         mkl  
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py39h7f8727e_0  
[conda] mkl_fft                   1.3.1            py39hd3c417c_0  
[conda] mkl_random                1.2.2            py39h51133e4_0  
[conda] numpy                     1.23.3                   pypi_0    pypi
[conda] numpy-base                1.23.1           py39ha15fc14_0  
[conda] pytorch                   1.13.0.dev20221001 py3.9_cuda11.7_cudnn8.5.0_0    pytorch-nightly
[conda] pytorch-cuda              11.7                 h67b0de4_0    pytorch-nightly
[conda] pytorch-mutex             1.0                        cuda    pytorch-nightly
[conda] torch                     1.13.0.dev20220923+cu116          pypi_0    pypi
[conda] torchaudio                0.13.0.dev20220923+cu116          pypi_0    pypi
[conda] torchvision               0.14.0.dev20220923+cu116          pypi_0    pypi
```

cc @ezyang @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7 @Chillee @samdow @soumith