RuntimeError: Only tuples, lists and Variables supported as JIT inputs/outputs. Dictionaries and strings are also accepted but their usage is not recommended.

Hi, I have issue with exporting pytorch models to onnx with output of List[Dict[str, tensor]], just like some detection models in torchvision.

### test return with dict
import torch
import onnx
import io
import onnx.helper as helper
from torch.jit.annotations import Tuple, List, Dict, Optional
from torch import Tensor

class DummyModel(torch.nn.Module):
    def __init__(self):
        super(DummyModel, self).__init__()

    def forward(self, x: Tensor, y: Tensor, z: Tensor)->List[Dict[str, Tensor]]:
        res = torch.jit.annotate(List[Dict[str, torch.Tensor]], [])
        xy = x + y
        xz = x + z
        yz = y + z
        res.append({'xy' : xy, 'xz' : xz, 'yz' : yz})
        return res

model = DummyModel()
x = torch.arange(4).reshape(2,2)
y = torch.arange(4).reshape(2,2)
z = torch.arange(4).reshape(2,2)
input_data = (x, y, z)

desired = model(*input_data)
print(f'Expect:\n{desired}')

# onnx_io = "test_dict.onnx"
torch.onnx.export(
    model,
    input_data,
    onnx_io,
    verbose=False,
    input_names=None,
    output_names=None,
)

model_onnx = onnx.load(onnx_io)

sess = rt.InferenceSession(onnx_io)
input_names = [_.name for _ in sess.get_inputs()]
print(f'input_names={input_names}')
#forward model
input_data_npy = [_.detach().cpu().numpy() for _ in input_data]
actual = sess.run(None, {name : data for name, data in zip(input_names, input_data_npy)})
desired = desired[0]
print(f'Actual:\n{actual}')

Any ideas? Thanks in advance.

cc @neginraoof