Repository: pytorch
Issue: Cannot export FeaStConv-based network to onnx via torch.jit.script
Labels: module: onnx,triaged

### 🐛 Describe the bug

 Hello!
 I planned to export a network based on graph neural network as an onnx model, but it always failed. Then I found that exporting a single-layer network `torch_geometric.nn.FeaStConv`  to onnx would also fail. I don’t know if any of you have encountered similar problems. I would like to know why and how to solve them. 
 Thanks!

what I run:
```python
import torch.onnx
from torch_geometric.nn import FeaStConv
from torch_geometric.data import Data


class GNNModule1(torch.nn.Module):
    def __init__(self, in_channel=3):
        super(GNNModule1, self).__init__()
        self.l_conv1 = FeaStConv(in_channel, 32, 9).jittable()

    def forward(self, x, edge_index, edge_weight):
        x = self.l_conv1(x, edge_index)
        return x

if __name__ =="__main__":


    num_nodes = 4
    in_channels = 3
    x = torch.randn((num_nodes, in_channels))
    edge_index = torch.tensor([
        [0, 1, 2, 3],  # 边的起始节点
        [1, 2, 3, 0]   # 边的终止节点
    ], dtype=torch.long)
    edge_weight = torch.tensor([1.0, 0.5, 1.2, 0.8], dtype=torch.float)
    data = Data(x=x, edge_index=edge_index, edge_weight=edge_weight)

    model = GNNModule1(in_channel=in_channels)
    model.eval()

    # output = model(data.x, data.edge_index, data.edge_weight)
    # print("Output:", output)

    model = torch.jit.script(model)   # comment this line will be success

    onnx_path = "gnn_model.onnx"
    try:
        torch.onnx.export(
            model,  # 脚本化后的模型
            (data.x, data.edge_index, data.edge_weight),  # 模型输入
            onnx_path,  # 输出文件路径
            input_names=['x', 'edge_index', 'edge_weight'],  # 输入的名称
            output_names=['out'],  # 输出的名称
            verbose=True,
            opset_version=15  # ONNX opset 版本
        )
        print(f"ONNX model saved to {onnx_path}")
    except Exception as e:
        print("Error during ONNX export:", e)
```

what I got:
```shell
[W C:\actions-runner\_work\pytorch\pytorch\builder\windows\pytorch\torch\csrc\jit\passes\lower_tuples.cpp:254] Warning: tuple appears in the op outputs, but this op does not forward tuples, unsupported kind: prim::unchecked_cast (function flattenOutputs)
Torch IR graph at exception: graph(%x.1 : Float(4, 3, strides=[3, 1], requires_grad=0, device=cpu),
      %edge_index.1 : Long(2, 4, strides=[4, 1], requires_grad=0, device=cpu),
      %edge_weight : Float(4, strides=[1], requires_grad=0, device=cpu),
      %l_conv1.u.weight : Float(9, 3, strides=[3, 1], requires_grad=0, device=cpu),
      %l_conv1.c : Float(9, strides=[1], requires_grad=0, device=cpu),
      %l_conv1.lin.weight : Float(288, 3, strides=[3, 1], requires_grad=0, device=cpu),
      %l_conv1.bias : Float(32, strides=[1], requires_grad=0, device=cpu)):
    ......
      %548 : int = aten::dim(%out.9), scope: GNNModule1::/torch_geometric.nn.aggr.basic.MeanAggregation::aggr_module # D:\ProgramData\anaconda3\envs\Mesh_Net\lib\site-packages\torch_geometric\utils\scatter.py:19:21
      %size.17 : int[] = aten::mul(%10, %548), scope: GNNModule1::/torch_geometric.nn.aggr.basic.MeanAggregation::aggr_module # D:\ProgramData\anaconda3\envs\Mesh_Net\lib\site-packages\torch_geometric\utils\scatter.py:19:15
      %550 : int[] = aten::_set_item(%size.17, %514, %19), scope: GNNModule1::/torch_geometric.nn.aggr.basic.MeanAggregation::aggr_module # D:\ProgramData\anaconda3\envs\Mesh_Net\lib\site-packages\torch_geometric\utils\scatter.py:20:8
      %551 : Tensor = aten::view(%count.13, %size.17), scope: GNNModule1::/torch_geometric.nn.aggr.basic.MeanAggregation::aggr_module # D:\ProgramData\anaconda3\envs\Mesh_Net\lib\site-packages\torch_geometric\utils\scatter.py:21:15
      %552 : Tensor = aten::expand_as(%551, %out.9), scope: GNNModule1::/torch_geometric.nn.aggr.basic.MeanAggregation::aggr_module # D:\ProgramData\anaconda3\envs\Mesh_Net\lib\site-packages\torch_geometric\utils\scatter.py:21:15
      %553 : Tensor = aten::div(%out.9, %552), scope: GNNModule1::/torch_geometric.nn.aggr.basic.MeanAggregation::aggr_module # D:\ProgramData\anaconda3\envs\Mesh_Net\lib\site-packages\torch_geometric\utils\scatter.py:84:19
      -> (%553)
  %bias.3 : Tensor = prim::GetAttr[name="bias"](%l_conv1), scope: GNNModule1::
  %x.5 : Tensor = aten::add(%out.7, %bias.3, %31), scope: GNNModule1:: # C:\Users\WooChi\AppData\Local\Temp\WooChi_pyg\tmpc6s_oinc.py:247:18
  return (%x.5)

Error during ONNX export: tuple use not matched to tuple construct. Instead found: prim::Return
```
It seems to be an error caused by outputting  a tuple (the log is very long and complicated, forgive me for not understanding it), but the output of the network is actually a tensor, and when `torch.jit.script` is not used, it can be exported normally as onnx model.
 
(Actually, I searched issues and found #41548 which is similar but seems different.)

### Versions

My environment:
```shell
python collect_env.py
Collecting environment information...
PyTorch version: 1.13.0+cu116
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 11 专业版 (10.0.22631 64 位)
GCC version: Could not collect
Clang version: Could not collect
CMake version: version 3.30.2
Libc version: N/A

Python version: 3.8.16 (default, Jun 12 2023, 21:00:42) [MSC v.1916 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.22631-SP0
Is CUDA available: True
CUDA runtime version: 11.6.55
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA RTX 2000 Ada Generation Laptop GPU
Nvidia driver version: 556.12
cuDNN version: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.6\bin\cudnn_ops_train64_8.dll
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Name: 13th Gen Intel(R) Core(TM) i7-13850HX
Manufacturer: GenuineIntel
Family: 198
Architecture: 9
ProcessorType: 3
DeviceID: CPU0
CurrentClockSpeed: 2100
MaxClockSpeed: 2100
L2CacheSize: 28672
L2CacheSpeed: None
Revision: None

Versions of relevant libraries:
[pip3] numpy==1.24.4
[pip3] onnx==1.16.1
[pip3] torch==1.13.0+cu116
[pip3] torch-cluster==1.6.1+pt113cu116
[pip3] torch-geometric==2.3.1
[pip3] torch-scatter==2.1.1+pt113cu116
[pip3] torch-sparse==0.6.17+pt113cu116
[pip3] torchaudio==0.13.0+cu116
[pip3] torchvision==0.14.0+cu116
[conda] numpy                     1.24.4                   pypi_0    pypi
[conda] torch                     1.13.0+cu116             pypi_0    pypi
[conda] torch-cluster             1.6.1+pt113cu116          pypi_0    pypi
[conda] torch-geometric           2.3.1                    pypi_0    pypi
[conda] torch-scatter             2.1.1+pt113cu116          pypi_0    pypi
[conda] torch-sparse              0.6.17+pt113cu116          pypi_0    pypi
[conda] torchaudio                0.13.0+cu116             pypi_0    pypi
[conda] torchvision               0.14.0+cu116             pypi_0    pypi
```