Cannot run scripted BERT_Pytorch
Labels: oncall: jit

### 🐛 Describe the bug

I tried to run the following code, but it fails with error


Traceback (most recent call last):
  File "bug.py", line 70, in <module>
    main()
  File "bug.py", line 65, in main
    script_f(*arg_list)
  File ".../pytorch/torch/nn/modules/module.py", line 1131, 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)

```python
import torch        
from functorch import make_fx
from torch.nn.utils.stateless import functional_call
from functorch.compile import ts_compile
from functorch.compile import default_decompositions
import torch.utils._pytree as pytree
import importlib
import gc

def load_model(device, model_name):
    module = importlib.import_module(f"torchbenchmark.models.{model_name}")
    benchmark_cls = getattr(module, "Model", None)
    if not hasattr(benchmark_cls, "name"):
        benchmark_cls.name = model_name
    batch_size = None
    
    benchmark = benchmark_cls(
        test="train", device=device, jit=False, batch_size=batch_size
    )

    model, example_inputs = benchmark.get_module()

    model.eval()
    gc.collect()
    return device, benchmark.name, model, example_inputs


def trace_model(model, inputs):
    def f(params, inp):
        out = functional_call(model, params, inp)
        # out.sum().backward()
        result = 0
        if isinstance(out, tuple):
            for i in out:
                result += i.sum()
        else:
            result = out.sum()
        result.sum().backward()
        return [param.grad for param in params.values()]
    
    params = dict(model.named_parameters())
    traced_graph = make_fx(f, decomposition_table=default_decompositions)(params, inputs)
    return traced_graph, params

def main():
    torch._C._jit_override_can_fuse_on_cpu(False)
    torch._C._jit_override_can_fuse_on_gpu(False)
    torch._C._jit_set_texpr_fuser_enabled(False)
    torch._C._jit_set_nvfuser_enabled(True)


    torch.manual_seed(1337)
    device, name, model, example_inputs = load_model(
        "cuda", 'BERT_pytorch'
    )


    traced_graph, params = trace_model(model, example_inputs)
    traced_graph.graph.set_codegen(torch.fx.graph.CodeGen())  # avoid recursive pytree
    script_f = ts_compile(traced_graph, (params, example_inputs))
    arg_list, spec  = pytree.tree_flatten([params, example_inputs])
    script_f(*arg_list)
    script_f(*arg_list)
    script_f(*arg_list)
    script_f(*arg_list)


if __name__ == "__main__":
    main()
```

### Versions

Collecting environment information...
PyTorch version: 1.13.0a0+git9244547
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (conda-forge gcc 12.1.0-16) 12.1.0
Clang version: 6.0.0-1ubuntu2 (tags/RELEASE_600/final)
CMake version: version 3.22.5
Libc version: glibc-2.27

Python version: 3.8.13 (default, Mar 28 2022, 11:38:47)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-1069-aws-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.6.112
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB
GPU 4: NVIDIA A100-SXM4-40GB
GPU 5: NVIDIA A100-SXM4-40GB
GPU 6: NVIDIA A100-SXM4-40GB
GPU 7: NVIDIA A100-SXM4-40GB

Nvidia driver version: 510.47.03
cuDNN version: Probably one of the following:
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.1.1
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: False

Versions of relevant libraries:
[pip3] bert-pytorch==0.0.1a4
[pip3] botorch==0.6.4
[pip3] flowtorch==0.8
[pip3] functorch==0.3.0a0+6cfb462
[pip3] gpytorch==1.6.0
[pip3] mypy==0.950
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.3
[pip3] numpyro==0.9.2
[pip3] pytorch-transformers==1.2.0
[pip3] torch==1.13.0a0+git9244547
[pip3] torch-struct==0.5
[pip3] torchaudio==0.12.0a0+4d2fa19
[pip3] torchdynamo==0.2.0
[pip3] torchfile==0.1.0
[pip3] torchmetrics==0.9.1
[pip3] torchrec-nightly==2022.4.26
[pip3] torchtext==0.13.0a0+d6e3550
[pip3] torchvision==0.14.0a0+a7e4fbd
[pip3] torchx-nightly==2022.6.8
[conda] bert-pytorch              0.0.1a4                   dev_0    <develop>
[conda] blas                      1.0                         mkl  
[conda] botorch                   0.6.4                    pypi_0    pypi
[conda] flowtorch                 0.8                      pypi_0    pypi
[conda] functorch                 0.3.0a0+6cfb462           dev_0    <develop>
[conda] gpytorch                  1.6.0                    pypi_0    pypi
[conda] magma-cuda110             2.5.2                         1    pytorch
[conda] magma-cuda116             2.6.1                         0    pytorch
[conda] mkl                       2022.1.0                 pypi_0    pypi
[conda] mkl-include               2022.0.1           h06a4308_117  
[conda] mkl-service               2.4.0            py38h7f8727e_0  
[conda] mkl_fft                   1.3.1            py38hd3c417c_0  
[conda] mkl_random                1.2.2            py38h51133e4_0  
[conda] numpy                     1.20.0                   pypi_0    pypi
[conda] numpy-base                1.22.3           py38hf524024_0  
[conda] numpyro                   0.9.2                    pypi_0    pypi
[conda] pytorch-transformers      1.2.0                    pypi_0    pypi
[conda] torch                     1.13.0a0+git9244547           dev_0    <develop>
[conda] torch-struct              0.5                      pypi_0    pypi
[conda] torchaudio                0.12.0a0+4d2fa19           dev_0    <develop>
[conda] torchdynamo               0.2.0                     dev_0    <develop>
[conda] torchfile                 0.1.0                    pypi_0    pypi
[conda] torchmetrics              0.9.0                    pypi_0    pypi
[conda] torchrec-nightly          2022.4.26                pypi_0    pypi
[conda] torchtext                 0.13.0a0+d6e3550           dev_0    <develop>
[conda] torchvision               0.14.0a0+a7e4fbd           dev_0    <develop>
[conda] torchx-nightly            2022.6.6                 pypi_0    pypi