Repository: pytorch
Issue: [JIT] Consecutive use of `addmm` Leads to Exception
Labels: oncall: jit

### 🐛 Describe the bug

The following program behave inconsistently between eager mode and JIT, which arises from consecutive use of `addmm`. If `torch.jit.optimize_for_inference` is removed, the program works fine.
To reproduce:
```python
import torch

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.v1 = torch.nn.Parameter(torch.rand([64, 32]))
        self.v2 = torch.nn.Parameter(torch.rand([64, 32]))
        self.v3 = torch.nn.Parameter(torch.rand([32, 32]))

    def forward(self, i1, i2):
        r1 = torch.addmm(self.v2, self.v1, i1)
        r2 = torch.addmm(r1, i2, self.v3)
        return r2


i1 = torch.rand([32, 32])
i2 = torch.rand([64, 32])

fn = Model()
fn.eval()

fn(i1, i2)
print("Eager model ... OK!")

exported = torch.jit.trace(fn, (i1, i2))
exported = torch.jit.optimize_for_inference(exported) # works w/o this line
exported(i1, i2)
print("JIT ... OK!")

"""
Eager model ... OK!
Traceback (most recent call last):
  File "model.py", line 27, in <module>
    exported(i1, i2)
  File "python3.10/site-packages/torch/nn/modules/module.py", line 1488, in _call_impl
    return forward_call(*args, **kwargs)
RuntimeError: 
The size of tensor a (64) must match the size of tensor b (32) at non-singleton dimension 0
The above operation failed shape propagation in this context:
  File "<string>", line 3

      def addmm(self: Tensor, mat1: Tensor, mat2: Tensor, beta: number = 1.0, alpha: number = 1.0):
          return self + mat1.mm(mat2)
                 ~~~~~~~~~~~~~~~~~~~ <--- HERE

      def batch_norm(input : Tensor, running_mean : Optional[Tensor], running_var : Optional[Tensor], training : bool, momentum : float, eps : float) -> Tensor:
"""
```

### Versions

<details><summary><b>Environment</b> <i>[Click to expand]</i></summary>

```
PyTorch version: 2.0.0.dev20230119+cu117
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.3.0-1ubuntu1~22.04) 11.3.0
Clang version: 11.1.0-6
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.10.6 (main, Nov 14 2022, 16:10:14) [GCC 11.3.0] (64-bit runtime)
Python platform: Linux-5.15.0-56-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 11.6.124
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090
GPU 2: NVIDIA GeForce RTX 3090

Nvidia driver version: 515.86.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.4.1
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.5
[pip3] pytorch-triton==2.0.0+0d7e753227
[pip3] torch==2.0.0.dev20230119+cu117
[pip3] torchaudio==2.0.0.dev20230118+cu117
[pip3] torchvision==0.15.0.dev20230118+cu117
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.3.1               h2bc3f7f_2  
[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.21.5           py39he7a7128_1  
[conda] numpy-base                1.21.5           py39hf524024_1  
[conda] numpydoc                  1.2                pyhd3eb1b0_0  
[conda] torch                     1.14.0a0+gitce2f870          pypi_0    pypi
```

</details>

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel