How to export GNN with dict inputs correctly?
Labels: module: onnx, triaged

## Problem description

I am having an issue when exporting of PyTorch GNN model to ONNX. Here is my export code:

```
torch.onnx.export(
    model=model,
    args=(x_dict, edge_index_dict, edge_attr_dict, {}),
    f=save_path,
    verbose=False,
    input_names=["x_dict", "edge_index_dict", "edge_attr_dict"],
    output_names=["out"],
)
```

`x_dict, edge_index_dict, edge_attr_dict` are of type `Dict[str, torch.Tensor]` (hetero_data is formed [like this](https://github.com/emnigma/VSharp/blob/408ba9800362285f420b3d9b51116f4b2cbb3391/VSharp.ML.AIAgent/ml/data_loader_compact.py#L30))

In addition to 3 inputs in my [model](https://github.com/emnigma/VSharp/blob/408ba9800362285f420b3d9b51116f4b2cbb3391/VSharp.ML.AIAgent/ml/models.py#L654)'s [forward](https://github.com/emnigma/VSharp/blob/408ba9800362285f420b3d9b51116f4b2cbb3391/VSharp.ML.AIAgent/ml/models.py#L659) , torch.onnx.export generates 4 additional inputs and when I try to use exported model with onnxruntime I get ValueError:

`ValueError: Required inputs (['edge_index', 'edge_index.5', 'edge_index.3', 'onnx::Reshape_9']) are missing from input feed (['x_dict', 'edge_index_dict', 'edge_attr_dict']).`

I am getting a feeling I am doing something wrong, how can i export my model correctly?

## Reproduction

here is a minimal reproduction script and dummy_data for it:

script: https://gist.github.com/emnigma/0b98cfbf3fff47be417c64489d83a2a2

data: https://gist.github.com/emnigma/e3ea559fe4db0adde886708f402473bb

## JIT trace output

I also tried to trace model compilation, here is the jit trace results with strict=False .code output:

```
def forward(self,
    argument_1: Dict[str, Tensor],
    argument_2: Dict[str, Tensor],
    argument_3: Dict[str, Tensor]) -> Dict[str, Tensor]:
  state_encoder = self.state_encoder
  x = argument_1["game_vertex"]
  x0 = argument_1["state_vertex"]
  edge_index = argument_2["game_vertex to game_vertex"]
  edge_index0 = argument_2["game_vertex in state_vertex"]
  edge_index1 = argument_2["game_vertex history state_vertex"]
  edge_index2 = argument_2["state_vertex parent_of state_vertex"]
  edge_weight = argument_3["game_vertex history state_vertex"]
  _0 = (state_encoder).forward(x, edge_index, x0, edge_index2, edge_index1, edge_weight, edge_index0, )
  _1 = {"state_vertex": _0, "game_vertex": x}
  return _1
```

## System Info
PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.4.1 (arm64)
GCC version: Could not collect
Clang version: 14.0.3 (clang-1403.0.22.14.1)
CMake version: version 3.26.4
Libc version: N/A

Python version: 3.11.4 (main, Jul  5 2023, 08:40:20) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-13.4.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.25.0
[pip3] torch==2.0.1
[pip3] torch-geometric==2.3.1
[pip3] torch-scatter==2.1.1
[pip3] torch-sparse==0.6.17
[pip3] torchaudio==2.0.2
[pip3] torchvision==0.15.2a0
[conda] numpy                     1.25.0          py311he598dae_0  
[conda] numpy-base                1.25.0          py311hfbfe69c_0  
[conda] pytorch                   2.0.1                  py3.11_0    pytorch
[conda] torch-geometric           2.3.1                    pypi_0    pypi
[conda] torch-scatter             2.1.1                    pypi_0    pypi
[conda] torch-sparse              0.6.17                   pypi_0    pypi
[conda] torchaudio                2.0.2                 py311_cpu    pytorch
[conda] torchvision               0.15.2          cpu_py311he74fb5d_0