Repository: pytorch
Issue: TorchScript causes range_check error after a few iterations of forward-backward passes
Labels: oncall: jit

### 🐛 Describe the bug

I came across an error caused by a TorchScript model after two iterations of forward-backward passes.
The following script can reproduce the error.
The serialized model can be downloaded from: [repro_model.bin.zip](https://github.com/pytorch/pytorch/files/9662120/repro_model.bin.zip)

```python
import torch

m = torch.jit.load("repro_model.bin")
print("Graph: {}".format(m.graph))

x1 = torch.randn(1, 300, 1024, requires_grad=True).cuda()
x2 = torch.randn(1, 300, 1024, requires_grad=True).cuda()

for i in range(0, 3):
    y = m((x1, x2))
    tgt = torch.randn_like(y[0])
    y[0].backward(tgt)
    print("Iteration {} finished".format(i))
```

The output of the script is shown below.
Note that the first two iterations successfully finished, and the third iteration produced an error.

The state of the model should be the same at the beginning of every iteration.
Once a pair of forward/backward passes is successfully finished, it is expected that the following iteration also succeeds.

```bash
Graph: graph(%self : __torch__.test_module,
      %merged_input.1 : (Tensor, Tensor)):
  %46 : Tensor = prim::Constant[value={1e-12}]() # :0:0
  %23 : bool = prim::Constant[value=1]() # :0:0
  %17 : NoneType = prim::Constant()
  %11 : int = prim::Constant[value=1]()
  %7 : bool = prim::Constant[value=0]() # :0:0
  %6 : float = prim::Constant[value=0.10000000000000001]() # :0:0
  %14 : int = prim::Constant[value=6]() # :0:0
  %19 : int = prim::Constant[value=-1]() # :0:0
  %36 : int = prim::Constant[value=2]() # :0:0
  %_input.1 : Tensor, %_input_tensor.1 : Tensor = prim::TupleUnpack(%merged_input.1)
  %_hidden_states.1 : Tensor = aten::dropout(%_input.1, %6, %7) # :0:0
  %_x.1 : Tensor = aten::add(%_hidden_states.1, %_input_tensor.1, %11) # :0:0
  %_x0.1 : Tensor = aten::to(%_x.1, %14, %7, %7, %17) # :0:0
  %20 : int[] = prim::ListConstruct(%19)
  %_u.1 : Tensor = aten::mean(%_x0.1, %20, %23, %17) # :0:0
  %30 : Tensor = aten::sub(%_x0.1, %_u.1, %11) # :0:0
  %34 : Tensor = aten::sub(%_x0.1, %_u.1, %11) # :0:0
  %37 : Tensor = aten::pow(%34, %36) # :0:0
  %38 : int[] = prim::ListConstruct(%19)
  %_s.1 : Tensor = aten::mean(%37, %38, %23, %17) # :0:0
  %48 : Tensor = aten::add(%_s.1, %46, %11) # :0:0
  %50 : Tensor = aten::sqrt(%48) # :0:0
  %_x1.1 : Tensor = aten::div(%30, %50) # :0:0
  %55 : (Tensor) = prim::TupleConstruct(%_x1.1)
  return (%55)

Iteration 0 finished
Iteration 1 finished
Traceback (most recent call last):
  File "run_graph_test.py", line 10, in <module>
    y = m((x1, x2))
  File "/.../python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
RuntimeError: vector::_M_range_check: __n (which is 18446744073709551615) >= this->size() (which is 3)
```





### Versions

I used PyTorch 1.12.0 with CUDA 11.3 on NVIDIA A100.


```bash
$ python collect_env.py
Collecting environment information...
PyTorch version: 1.12.0+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Red Hat Enterprise Linux release 8.4 (Ootpa) (x86_64)
GCC version: (GCC) 8.4.1 20200928 (Red Hat 8.4.1-1)
Clang version: 11.0.1 (Red Hat 11.0.1-1.module+el8.4.0+12483+89b287b0)
CMake version: version 3.24.1
Libc version: glibc-2.28

Python version: 3.8.14 (default, Sep 16 2022, 16:31:55)  [GCC 8.4.1 20200928 (Red Hat 8.4.1-1)] (64-bit runtime)
Python platform: Linux-4.18.0-305.25.1.el8_4.x86_64-x86_64-with-glibc2.28
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration:
GPU 0: A100-SXM-80GB
GPU 1: A100-SXM-80GB
GPU 2: A100-SXM-80GB
GPU 3: A100-SXM-80GB
GPU 4: A100-SXM-80GB
GPU 5: A100-SXM-80GB
GPU 6: A100-SXM-80GB
GPU 7: A100-SXM-80GB

Nvidia driver version: 460.73.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.12.0+cu113
[conda] No relevant packages
```