Repository: pytorch
Issue: `torch.cat` can break `torch.jit.ScriptModule` when in inference mode
Labels: oncall: jit

### 🐛 Describe the bug

Given a `ScriptModule` that concatenates two tensors and uses the output in another op, it breaks under `inference_mode`.

Code to reproduce:
```python
import torch


class Model(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.x = torch.nn.Parameter(data=torch.tensor(0.0))

    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        return self.x * torch.cat([a, b], dim=1)

a = torch.zeros(1, 1)
b = torch.zeros(1, 1)
model = torch.jit.script(Model())
model(a, b)  # succeeds
with torch.inference_mode():
    model(a, b)  # fails
```

Traceback:
```
self = Model(
  (linear): RecursiveScriptModule(original_name=Linear)
), input = (tensor([[0.]]), tensor([[0.]])), kwargs = {}, forward_call = <torch.ScriptMethod object at 0x114de8130>

    def _call_impl(self, *input, **kwargs):
        forward_call = (self._slow_forward if torch._C._get_tracing_state() else self.forward)
        # If we don't have any hooks, we want to skip the rest of the logic in
        # this function, and just call forward.
        if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
                or _global_forward_hooks or _global_forward_pre_hooks):
>           return forward_call(*input, **kwargs)
E           RuntimeError: The following operation failed in the TorchScript interpreter.
E           Traceback of TorchScript (most recent call last):
E           RuntimeError: Inference tensors cannot be saved for backward. To work around you can make a clone to get a normal tensor and use it in autograd.

venv/lib/python3.8/site-packages/torch/nn/modules/module.py:1130: RuntimeError
```

Observations:

1. Running the model without inference mode succeeds
2. Running the model in inference mode without first running it *outside* of inference mode succeeds
3. Only after running the model outside of inference mode first, does a subsequent run from within inference mode fail
4. This behavior happens when the output of `torch.cat` is used in an op with `torch.nn.Parameter`s or trainable layers (e.g. `torch.nn.Linear`), but not with constants.
5. This behavior occurs whether converting the module to a `ScriptModule` after instantiation via `torch.jit.script` or if directly subclassing `torch.jit.ScriptModule`.
6. The model fails in the same way if first traced with `torch.jit.trace`.

### Versions

```
PyTorch version: 1.12.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.4 (x86_64)
GCC version: Could not collect
Clang version: 13.1.6 (clang-1316.0.21.2.5)
CMake version: version 3.23.2
Libc version: N/A

Python version: 3.8.13 (default, Apr 13 2022, 19:33:23)  [Clang 13.1.6 (clang-1316.0.21.2.3)] (64-bit runtime)
Python platform: macOS-12.4-x86_64-i386-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] efficientnet-pytorch==0.7.1
[pip3] mypy==0.971
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.6
[pip3] pytorch-ranger==0.1.1
[pip3] torch==1.12.0
[pip3] torch-optimizer==0.1.0
[pip3] torch-tb-profiler==0.4.0
[pip3] torchmetrics==0.9.3
[pip3] torchvision==0.13.0
[conda] Could not collect
```