Hi, how to export models to onnx with output of List[Dict[str, tensor]], just like detection models in torchvision.
Following code mocking here, but does not work.
Any ideas?
Thanks in advance.
### 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: int)->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()
input_data = (torch.arange(4).reshape(2,2) for _ in range(3))

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

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

model_onnx = onnx.load(onnx_io)

# print(helper.printable_graph(model_onnx.graph))
sess = rt.InferenceSession(out_path)
input_names = [_.name for _ in sess.get_inputs()]
print(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)})
actual = actual[0]
print(f'Actual:\n{actual}')
for k, v in desired.items():
    np.testing.assert_allclose(desired[k], actual[k])