[ONNX] `.squeeze(1)` on the B X T (not B X 1 X T) tensor causes export error in masking
Labels: module: onnx, triaged, onnx-triaged

### 🐛 Describe the bug

Here is I made repro script with simple model

```python
import torch
import torch.nn as nn
import onnxruntime

device = 'cuda'
use_mask = True
model_path = 'simple_model.onnx' 


### Mask calculation

def compute_output_lengths(x, lengths_fraction=None):
    if lengths_fraction is None:
        return torch.full(x.shape[:1], x.shape[-1], device=x.device, dtype=torch.long)
    return (lengths_fraction * x.shape[-1]).ceil().long()


def temporal_mask(x, lengths):
    return (torch.arange(x.shape[-1], device=x.device, dtype=lengths.dtype).unsqueeze(0) <
            lengths.unsqueeze(1)).view(x.shape[:1] + (1,) * (len(x.shape) - 2) + x.shape[-1:])


### Simple model for export

class SimpleNetwork(nn.Module):
    def __init__(self, use_mask=False):
        super().__init__()
        self.use_mask = use_mask
        self.conv1 = nn.Conv1d(in_channels=1,
                               out_channels=3,
                               kernel_size=3,
                               stride=1,
                               padding=0,
                               dilation=1,
                               groups=1
                               )
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv1d(in_channels=3,
                               out_channels=1,
                               kernel_size=3,
                               stride=1,
                               padding=0,
                               dilation=1,
                               groups=1)

    def forward(self, x, xlen):
        # x - [B; T], xlen - [B]
        x = x.unsqueeze(1)
        # x - [B; 1; T]
        if self.use_mask:
            mask = temporal_mask(x, compute_output_lengths(x, xlen))
            x = x * mask    
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        return x


### Random tensor to export

onnx_sample_batch_size = 16
onnx_sample_time = 1024

waveform_input = torch.rand(onnx_sample_batch_size, onnx_sample_time, device=device)
xlen = torch.rand(onnx_sample_batch_size, device=device)

### Create model

model = SimpleNetwork(use_mask=use_mask).to(device)
result_torch = model(waveform_input, xlen)

### Export model

torch.onnx.export(
    model, (waveform_input, xlen,),
    model_path,
    verbose=False,
    opset_version=12,
    export_params=True,
    do_constant_folding=True,
    input_names=['x', 'xlen'],
    output_names=['logits'],
    dynamic_axes=dict(x={
        0: 'B', 1: 'T'
    }, logits={
        0: 'B', 2: 't'
    }, xlen={
        0: 'B'
    })
)

session = onnxruntime.InferenceSession(model_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
onnx_input = dict(x=waveform_input.cpu().numpy())
if use_mask:
    onnx_input['xlen'] = xlen.cpu().numpy()
result_onnx = session.run(None, onnx_input)[0]
result_onnx = torch.as_tensor(result_onnx, device=device)

### Сorrectness check

assert torch.allclose(result_torch.cpu(), result_onnx.cpu(), rtol=1e-02, atol=1e-03)

### Doing the same but with different shape

validate_batch_size = 32
validate_sample_time = 512

validate_waveform_input = torch.rand(validate_batch_size, validate_sample_time, device=device)
validate_xlen = torch.rand(validate_batch_size, device=device)

validate_result_torch = model(validate_waveform_input, validate_xlen)

validate_onnx_input = dict(x=validate_waveform_input.cpu().numpy())
if use_mask:
    validate_onnx_input['xlen'] = validate_xlen.cpu().numpy()
validate_result_onnx = session.run(None, validate_onnx_input)[0]
validate_result_onnx = torch.as_tensor(validate_result_onnx, device=device)

assert torch.allclose(validate_result_torch.cpu(), validate_result_onnx.cpu(), rtol=1e-02, atol=1e-03)
```

Here is fully convertible and working script.
But if i will add `.squeeze(1)` in forward like this:

```python
    def forward(self, x, xlen):
        # x - [B; T]
        # x.squeeze(1) - [B; T]
        # x.squeeze(1).unsqueeze(1) - [B; 1; T]
        x = x.squeeze(1).unsqueeze(1)
        if self.use_mask:
            mask = temporal_mask(x, compute_output_lengths(x, xlen))
            x = x * mask
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        return x
```

Then I'll get following exception:
```
[E:onnxruntime:, sequential_executor.cc:346 Execute] Non-zero status code returned while running Reshape node. Name:'Reshape_16' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/tensor/reshape_helper.h:41 onnxruntime::ReshapeHelper::ReshapeHelper(const onnxruntime::TensorShape&, std::vector<long int>&, bool) gsl::narrow_cast<int64_t>(input_shape.Size()) == size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{32,1024}, requested shape:{16,1,1024}
```

So `.squeeze(1)` on tensor without B X 1 X ... structure causes to static shapes fixing in the onnx model.
Also I found that this script works with squeeze if `mask=False`. So this bag is combination of mask calculation and squeeze.

Here is image from Netron. Problem in `Mul` where `B = 1024`

<img width="248" alt="image" src="https://user-images.githubusercontent.com/21354805/172596256-89ee7656-d7f2-4c59-895f-58244e768d05.png">

Another issue about it: https://github.com/pytorch/pytorch/issues/36796


### Versions

```
PyTorch version: 1.10.1
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.9.7 (default, Sep 16 2021, 13:09:58)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-3.10.0-1160.49.1.el7.x86_64-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.3.109
GPU models and configuration: 
GPU 0: Tesla V100-SXM2-32GB
GPU 1: Tesla V100-SXM2-32GB
GPU 2: Tesla V100-SXM2-32GB
GPU 3: Tesla V100-SXM2-32GB

Nvidia driver version: 470.82.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.2.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.2
[pip3] torch==1.10.1
[pip3] torch-audiomentations==0.9.1
[pip3] torch-pitch-shift==1.2.0
[pip3] torch-stft==0.1.4
[pip3] torchaudio==0.10.1
[pip3] torchvision==0.11.2
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.3.1               ha36c431_9    nvidia
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[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.2           py39h20f2e39_0  
[conda] numpy-base                1.21.2           py39h79a1101_0  
[conda] pytorch                   1.10.1          py3.9_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch-audiomentations     0.9.1                    pypi_0    pypi
[conda] torch-pitch-shift         1.2.0                    pypi_0    pypi
[conda] torch-stft                0.1.4                    pypi_0    pypi
[conda] torchaudio                0.10.1               py39_cu113    pytorch
[conda] torchvision               0.11.2               py39_cu113    pytorch
```

```
python -c 'import onnxruntime; print(onnxruntime.__version__)'
```
`1.10.0`