Repository: pytorch
Issue: jit.freeze throws RuntimeError: stack_out && stack_out->size() == 1 INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/frozen_conv_folding.cpp":281
Labels: oncall: jit

### 🐛 Describe the bug

When trying to `jit.freeze` my torchscript module, I am encountering an `INTERNAL ASSERT FAILED` error. I would have expected the freezing to simply run without errors. I have cut down my code to this minimal sample:

```python
import torch
from torch import nn

device = 'cuda'


def get_dummy_input():
    img_seq = torch.randn(3, 5, 3, 256, 256, device=device)
    return img_seq


class MinimalHead(nn.Module):
    def __init__(
        self,
        num_inp=5,
    ):
        super().__init__()
        self.num_inp = num_inp

        self.inp_heads = nn.ModuleList([
            nn.Conv2d(3, 3, 32)
            for _ in range(num_inp)
        ])

    def forward(self, inp_seq):
        inp_features = 0
        for inp_i in range(self.num_inp):
            inp_features = inp_features + self.inp_heads[inp_i](inp_seq[:, inp_i])
        return inp_features


def main():
    head = MinimalHead().to(device)
    head = torch.jit.trace(head, get_dummy_input()).eval()
    torch.jit.freeze(head)


if __name__ == '__main__':
    main()
```

Running the above script results in this error and traceback:
```
Traceback (most recent call last):
  File "minimal_freeze_error.py", line 39, in <module>
    main()
  File "minimal_freeze_error.py", line 35, in main
    torch.jit.freeze(head)
  File "[...]/torch/jit/_freeze.py", line 119, in freeze
    run_frozen_optimizations(out, optimize_numerics, preserved_methods)
  File "[...]/torch/jit/_freeze.py", line 167, in run_frozen_optimizations
    torch._C._jit_pass_optimize_frozen_graph(mod.graph, optimize_numerics)
RuntimeError: stack_out && stack_out->size() == 1 INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/frozen_conv_folding.cpp":281, please report a bug to PyTorch.
```
Since the error message suggests reporting this, that is what I am doing.

Interestingly, this same script running in my laptop instead of my server results in a smooth run.
Another workaround (not ideal due to existing jit checkpoints) I have currently found is to replace my model's `forward` with the following equivalent code:
```python
    def forward(self, inp_seq):
        inp_features = [
            self.inp_heads[inp_i](inp_seq[:, inp_i])
            for inp_i in range(self.num_inp)
        ]
        inp_features = torch.sum(
            torch.stack(
                inp_features,
                dim=0,
            ),
            dim=0,
        )
        return inp_features
```

### Versions

```
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: CentOS Linux 7 (Core) (x86_64)
GCC version: (GCC) 4.8.5 20150623 (Red Hat 4.8.5-44)
Clang version: Could not collect
CMake version: version 2.8.12.2
Libc version: glibc-2.17

Python version: 3.8.11 (default, Aug  3 2021, 15:09:35)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-3.10.0-1160.49.1.el7.x86_64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.6.112
GPU models and configuration:
GPU 0: Tesla P100-PCIE-16GB
GPU 1: Tesla P100-PCIE-16GB
GPU 2: Tesla P100-PCIE-16GB
GPU 3: Tesla P100-PCIE-16GB
GPU 4: Tesla P100-PCIE-16GB
GPU 5: Tesla P100-PCIE-16GB
GPU 6: Tesla P100-PCIE-16GB
GPU 7: Tesla P100-PCIE-16GB

Nvidia driver version: 510.47.03
cuDNN version: Could not collect
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-extensions==0.4.3
[pip3] numpy==1.19.5
[pip3] pytorch-lightning==1.4.9
[pip3] pytorch-metric-learning==0.9.99
[pip3] torch==1.12.0+cu113
[pip3] torch-fidelity==0.3.0
[pip3] torchaudio==0.12.0+cu113
[pip3] torchmetrics==0.5.1
[pip3] torchvision==0.13.0+cu113
[conda] efficientnet-pytorch      0.7.1                    pypi_0    pypi
[conda] numpy                     1.19.5                   pypi_0    pypi
[conda] pytorch-lightning         1.4.9                    pypi_0    pypi
[conda] pytorch-metric-learning   0.9.99                   pypi_0    pypi
[conda] torch                     1.12.0+cu113             pypi_0    pypi
[conda] torch-fidelity            0.3.0                    pypi_0    pypi
[conda] torchaudio                0.12.0+cu113             pypi_0    pypi
[conda] torchmetrics              0.5.1                    pypi_0    pypi
[conda] torchvision               0.13.0+cu113             pypi_0    pypi
```