jit.trace dost not support nested dict outputs
Labels: oncall: jit

### 🐛 Describe the bug

We have a complicated model which return value is a tuple of nested dict:
```python
def forward(self, inputs)
    # .....
    decoder_outputs = {'feature_a': inputs['feature_a'], 'feature_b': inputs['feature_b']}
    predictions = {
        'feature_a': {'targets': inputs['feature_a'], 'probs': inputs['feature_a']}, 
        'feature_b': {'targets': inputs['feature_b'], 'probs': inputs['feature_b']}
    }
    return predictions, decoder_outputs
```

We want to trace the model, save it as jit model type, but we got the following error:
```
of traced region cannot be understood by the tracer, only outputs matchingdict[Union[str, Tensor], Union[Tensor, Tuple[Tensor, ...]]] can be a dictionary output of a traced function
```

Fully reproducible scripts:

1) test on simple dict return
```python
import torch

inputs = {
    'feature_a': torch.Tensor([[101, 7770, 1107, 4638, 5381, 2605, 4511, 1351, 102, 0]]),
    'feature_b': torch.Tensor([[101, 7770, 1107, 4638, 5381, 2605, 4511, 1351, 102, 0]])
}

def fn_simple_dict(inputs):
    decoder_outputs = {'feature_a': inputs['feature_a'], 'feature_b': inputs['feature_b']}
    predictions = {
        'feature_a': inputs['feature_a'],
        'feature_b': inputs['feature_b']
    }
    return predictions, decoder_outputs

torch.jit.trace(fn_simple_dict, inputs, strict=False)
```

All is OK. `<torch.jit.ScriptFunction at 0x112096e50>`

2) test on nested dict return
```python
import torch

inputs = {
    'feature_a': torch.Tensor([[101, 7770, 1107, 4638, 5381, 2605, 4511, 1351, 102, 0]]),
    'feature_b': torch.Tensor([[101, 7770, 1107, 4638, 5381, 2605, 4511, 1351, 102, 0]])
}


def fn_nested_dict(inputs):
    decoder_outputs = {'feature_a': inputs['feature_a'], 'feature_b': inputs['feature_b']}
    predictions = {
        'feature_a': {'targets': inputs['feature_a'], 'probs': inputs['feature_a']},
        'feature_b': {'targets': inputs['feature_b'], 'probs': inputs['feature_b']}
    }
    return predictions, decoder_outputs

torch.jit.trace(fn_nested_dict, inputs, strict=False)
```

We got an error:
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Input In [6], in <cell line: 17>()
     11     predictions = {
     12         'feature_a': {'targets': inputs['feature_a'], 'probs': inputs['feature_a']},
     13         'feature_b': {'targets': inputs['feature_b'], 'probs': inputs['feature_b']}
     14     }
     15     return predictions, decoder_outputs
---> 17 torch.jit.trace(fn_nested_dict, inputs, strict=False)

File ~/works/tools/py38/.venv/lib/python3.8/site-packages/torch/jit/_trace.py:778, in trace(func, example_inputs, optimize, check_trace, check_inputs, check_tolerance, strict, _force_outplace, _module_class, _compilation_unit)
    772     raise AttributeError(
    773         "trace doesn't support compiling individual module's functions.\n"
    774         "Please use trace_module"
    775     )
    777 name = _qualified_name(func)
--> 778 traced = torch._C._create_function_from_trace(
    779     name, func, example_inputs, var_lookup_fn, strict, _force_outplace
    780 )
    782 # Check the trace against new traces created from user-specified inputs
    783 if check_trace:

RuntimeError: output 1 ({feature_a: {targets:   101  7770  1107  4638  5381  2605  4511  1351   102     0
[ CPUFloatType{1,10} ], probs:   101  7770  1107  4638  5381  2605  4511  1351   102     0
[ CPUFloatType{1,10} ]}, feature_b: {targets:   101  7770  1107  4638  5381  2605  4511  1351   102     0
[ CPUFloatType{1,10} ], probs:   101  7770  1107  4638  5381  2605  4511  1351   102     0
[ CPUFloatType{1,10} ]}}) of traced region cannot be understood by the tracer, only outputs matchingdict[Union[str, Tensor], Union[Tensor, Tuple[Tensor, ...]]] can be a dictionary output of a traced function
```


Is there any workaround for this error? Thank you a lot

### Versions

```
Collecting environment information...
PyTorch version: 1.8.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.3 (x86_64)
GCC version: Could not collect
Clang version: 13.0.0 (clang-1300.0.27.3)
CMake version: Could not collect

Python version: 3.8 (64-bit runtime)
Is CUDA available: False
CUDA runtime version: No CUDA
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

Versions of relevant libraries:
[pip] Could not collect
[conda] Could not collect
```

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel @VitalyFedyunin @mingfeima @XiaobingSuper @ashokei @jingxu10