ONNX-exported model cannot output Dict[str, X] or str
Labels: module: onnx, triaged

### 🐛 Describe the bug

`Dict[str, float]` or simply even `str` cannot be a return value from a scripted Pytorch model passed through `torch.onnx.export`

```
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import *

class ALSModel(nn.Module):
    def __init__(self, n_items: int, n_factors: int):
        super(ALSModel, self).__init__()

    def forward(self, ind: torch.Tensor, size: int = 5)-> str:
        return "heyo"
        
        
m = torch.jit.script(ALSModel(10000, 128))
x = (torch.ones(1).long(), 3)

torch.onnx.export(m, x, "data/model2.onnx", input_names=['contentId', 'size'],  verbose=True, opset_version=16)
```
Presents the error:
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In [1], line 17
     14 m = torch.jit.script(ALSModel(10000, 128))
     15 x = (torch.ones(1).long(), 3)
---> 17 torch.onnx.export(m, x, "data/model2.onnx", input_names=['contentId', 'size'],  verbose=True, opset_version=16)

File ~/micromamba/envs/onnx/lib/python3.10/site-packages/torch/onnx/__init__.py:350, in export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, custom_opsets, export_modules_as_functions)
     74 r"""
     75 Exports a model into ONNX format. If ``model`` is not a
     76 :class:`torch.jit.ScriptModule` nor a :class:`torch.jit.ScriptFunction`, this runs
   (...)
    345     model to the file ``f`` even if this is raised.
    346 """
    348 from torch.onnx import utils
--> 350 return utils.export(
    351     model,
    352     args,
    353     f,
    354     export_params,
    355     verbose,
    356     training,
    357     input_names,
    358     output_names,
    359     operator_export_type,
    360     opset_version,
    361     do_constant_folding,
    362     dynamic_axes,
    363     keep_initializers_as_inputs,
    364     custom_opsets,
    365     export_modules_as_functions,
    366 )

File ~/micromamba/envs/onnx/lib/python3.10/site-packages/torch/onnx/utils.py:163, in export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, custom_opsets, export_modules_as_functions)
    145 def export(
    146     model,
    147     args,
   (...)
    160     export_modules_as_functions=False,
    161 ):
--> 163     _export(
    164         model,
    165         args,
    166         f,
    167         export_params,
    168         verbose,
    169         training,
    170         input_names,
    171         output_names,
    172         operator_export_type=operator_export_type,
    173         opset_version=opset_version,
    174         do_constant_folding=do_constant_folding,
    175         dynamic_axes=dynamic_axes,
    176         keep_initializers_as_inputs=keep_initializers_as_inputs,
    177         custom_opsets=custom_opsets,
    178         export_modules_as_functions=export_modules_as_functions,
    179     )

File ~/micromamba/envs/onnx/lib/python3.10/site-packages/torch/onnx/utils.py:1074, in _export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, export_type, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, fixed_batch_size, custom_opsets, add_node_names, onnx_shape_inference, export_modules_as_functions)
   1071     dynamic_axes = {}
   1072 _validate_dynamic_axes(dynamic_axes, model, input_names, output_names)
-> 1074 graph, params_dict, torch_out = _model_to_graph(
   1075     model,
   1076     args,
   1077     verbose,
   1078     input_names,
   1079     output_names,
   1080     operator_export_type,
   1081     val_do_constant_folding,
   1082     fixed_batch_size=fixed_batch_size,
   1083     training=training,
   1084     dynamic_axes=dynamic_axes,
   1085 )
   1087 # TODO: Don't allocate a in-memory string for the protobuf
   1088 defer_weight_export = (
   1089     export_type is not torch.onnx.ExportTypes.PROTOBUF_FILE
   1090 )

File ~/micromamba/envs/onnx/lib/python3.10/site-packages/torch/onnx/utils.py:731, in _model_to_graph(model, args, verbose, input_names, output_names, operator_export_type, do_constant_folding, _disable_torch_constant_prop, fixed_batch_size, training, dynamic_axes)
    728 params_dict = _get_named_param_dict(graph, params)
    730 try:
--> 731     graph = _optimize_graph(
    732         graph,
    733         operator_export_type,
    734         _disable_torch_constant_prop=_disable_torch_constant_prop,
    735         fixed_batch_size=fixed_batch_size,
    736         params_dict=params_dict,
    737         dynamic_axes=dynamic_axes,
    738         input_names=input_names,
    739         module=module,
    740     )
    741 except Exception as e:
    742     torch.onnx.log("Torch IR graph at exception: ", graph)

File ~/micromamba/envs/onnx/lib/python3.10/site-packages/torch/onnx/utils.py:332, in _optimize_graph(graph, operator_export_type, _disable_torch_constant_prop, fixed_batch_size, params_dict, dynamic_axes, input_names, module)
    330 _C._jit_pass_lint(graph)
    331 if GLOBALS.onnx_shape_inference:
--> 332     _C._jit_pass_onnx_graph_shape_type_inference(
    333         graph, params_dict, GLOBALS.export_onnx_opset_version
    334     )
    335 return graph

RuntimeError: unexpected tensor scalar type
```

More realistically, I would like to output a `Dict[str, float]` which holds an external id and a score for that id. Seemingly this cannot be done, either because of ONNX itself, or the Pytorch export.

With Pytorch==1.12.1, this code simply spins forever, without it ever even hardcrashing for me.

```
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import *

class ALSModel(nn.Module):
    def __init__(self, n_items: int, n_factors: int):
        super(ALSModel, self).__init__()
        self.emb = nn.Embedding(n_items, n_factors)
        self.it2ind = {str(i): i for i in range(n_items)}# TODO use real values
        self.ind2it = {v:k for k,v in self.it2ind.items()}
        
    def map_out(self, scores: torch.Tensor, indices: torch.Tensor) -> Dict[str, float]:
        inds: List[int] = indices.squeeze().tolist()
        scorez: List[float] = scores.squeeze().tolist()
        return {self.ind2it[i]: s for i, s in zip(inds, scorez)}
        
    def forward(self, ind: torch.Tensor, size: int = 30)-> Dict[str, float]:
        u = self.emb(ind)
        scores = u @ self.emb.weight.t()
        s, i = scores.topk(size)
        return self.map_out(s, i)
    
m = torch.jit.script(ALSModel(10000, 128))
x = (torch.ones(1).long(), 3)

torch.onnx.export(m, x, "data/model2.onnx", input_names=['contentId', 'size'],  verbose=True, opset_version=16)
```

### Versions

PyTorch version: 1.12.1
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.1 LTS (x86_64)
GCC version: (Ubuntu 11.3.0-1ubuntu1~22.04) 11.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.10.6 | packaged by conda-forge | (main, Aug 22 2022, 20:36:39) [GCC 10.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-52-generic-x86_64-with-glibc2.35
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

Versions of relevant libraries:
[pip3] numpy==1.23.4
[pip3] torch==1.12.1
[pip3] torchvision==0.13.1
[conda] Could not collect   /// (micromamba==0.26.0)

This is related to #83801 , in which it was noted that `Dict`s might be on the roadmap. Is this the case? Because it's quite cumbersome to keep separate mapping files existing outside the model artifact itself, and I would really like to avoid it if possible.

Thanks for your time :slightly_smiling_face: