Repository: pytorch
Issue: Failing to convert torchaudio.models.Conformer to onnx after scripting of it via torch.jit.script
Labels: needs reproduction,module: onnx,triaged

### 🐛 Describe the bug

I create default conformer class (used params from docs), then script it:

```
conformer = torchaudio.models.Conformer(
    input_dim=80,
    num_heads=4,
    ffn_dim=128,
    num_layers=4,
    depthwise_conv_kernel_size=31)

traced_script_module = torch.jit.script(conformer)
```

After that onnx conversion fails:
```
torch.onnx.export(traced_script_module,
                  f="../model.onnx",
                  input_names=["input", "lengths"],
                  output_names=["hidden_states"],
                  args=(torch.ones(1, 1, 80), torch.ones(1, dtype=torch.long)),
                  dynamic_axes={
                    "input": { 0: "batch_count", 1: "batch_item_length" },
                    "lengths": { 0: "batch_count" },
                    "hidden_states": { 0: "batch_count", 1: "batch_item_length" }
                  }
                )
```

with:

> Conda\envs\torch\Lib\site-packages\torch\onnx\symbolic_opset9.py", line 7089, in prim_if
>     torch._C._jit_pass_onnx_node_shape_type_inference(
> RuntimeError: ScalarType UNKNOWN_SCALAR is an unexpected tensor scalar type

I'm not sure scripting is required for proper conformer conversion (is it?), however, I get the same error for more complex trained model which uses the same conformer class and may require scripting.



### Versions

> 
> PyTorch version: 2.2.1
> Is debug build: False
> CUDA used to build PyTorch: 12.1
> ROCM used to build PyTorch: N/A
> 
> OS: Майкрософт Windows 10 Pro
> GCC version: Could not collect
> Clang version: Could not collect
> CMake version: version 3.28.3
> Libc version: N/A
> 
> Python version: 3.12.2 | packaged by conda-forge | (main, Feb 16 2024, 20:42:31) [MSC v.1937 64 bit (AMD64)] (64-bit runtime)
> Python platform: Windows-10-10.0.19044-SP0
> Is CUDA available: True
> CUDA runtime version: Could not collect
> CUDA_MODULE_LOADING set to: LAZY
> GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3060
> Nvidia driver version: 551.61
> cuDNN version: Could not collect
> HIP runtime version: N/A
> MIOpen runtime version: N/A
> Is XNNPACK available: True
> 
> CPU:
> Architecture=9
> CurrentClockSpeed=3593
> DeviceID=CPU0
> Family=107
> L2CacheSize=4096
> L2CacheSpeed=
> Manufacturer=AuthenticAMD
> MaxClockSpeed=3593
> Name=AMD Ryzen 7 3700X 8-Core Processor
> ProcessorType=3
> Revision=28928
> 
> Versions of relevant libraries:
> [pip3] numpy==1.26.4
> [pip3] onnx==1.15.0
> [pip3] onnxruntime==1.17.1
> [pip3] onnxscript==0.1.0.dev20240315
> [pip3] torch==2.2.1
> [pip3] torchaudio==2.2.1
> [pip3] torchvision==0.17.1
> [conda] blas                      1.0                         mkl
> [conda] mkl                       2023.1.0         h6b88ed4_46358
> [conda] mkl-service               2.4.0           py312h2bbff1b_1
> [conda] mkl_fft                   1.3.8           py312h2bbff1b_0
> [conda] mkl_random                1.2.4           py312h59b6b97_0
> [conda] numpy                     1.26.4          py312hfd52020_0
> [conda] numpy-base                1.26.4          py312h4dde369_0
> [conda] pytorch                   2.2.1           py3.12_cuda12.1_cudnn8_0    pytorch
> [conda] pytorch-cuda              12.1                 hde6ce7c_5    pytorch
> [conda] pytorch-mutex             1.0                        cuda    pytorch
> [conda] torchaudio                2.2.1                    pypi_0    pypi
> [conda] torchvision               0.17.1                   pypi_0    pypi