[Question] Torch JIT/ONNX and Neo compilation

I'm attempting to get a fastai model optimized with SageMaker Neo compilation. It looks like I can save out my model in one of three ways.

- regular: `learn.save` which outputs e.g. `resnet34.pth`
- onnx: `torch.onnx.export` which outputs e.g. `resnet34.onnx`
- jit: `torch.jit.save` which outputs e.g. `resnet34.pth`

I've actually done all three to try and experiment with different ways of compiling the model with Neo. After saving in one of these ways I then attempt to compile the model. I've tried framework as pytorch and onnx with versions 1.0.0, and 1.4, respectively.

```python
framework='pytorch' # 'onnx'
framework_version='1.0.0' # '1.4'
output_path = '/'.join(estimator.output_path.split('/')[:-1])
optimized_ic = estimator.compile_model(target_instance_family='ml_c5',
input_shape={'actual_input_1':[1, 3, 224, 224]},
output_path=output_path,
framework= framework,
framework_version=framework_version,
role=role)
```

I get the same error trying all three ways, with both framework types:

> Failed Reason: Client Error: InputConfiguration: Invalid PyTorch model or input-shape mismatch. Make sure that inputs are lexically ordered and of the correct dimensionality. Upgrade to PyTorch 1.0 and use torch.jit.trace() and torch.jit.save() to export your model. Also cast your model with eval() before exporting.

Any idea what I might be doing wrong? And does the key name for `input_shape` matter if I save the model with `jit.save`?

------
### Details

For JIT I am doing this:

```python
learn.model.eval()
trace_input = torch.ones(1,3,args.image_size,args.image_size).cuda()
jit_model = torch.jit.trace(learn.model.float(), trace_input)
jit_model.eval() # tried doing eval in both places
output_path = str(path/f'{args.model_arch}_jit.pth')
torch.jit.save(jit_model, output_path)
```

And for ONNX I am doing this:

```python
dummy_input = torch.ones(1,3,args.image_size,args.image_size).cuda()
output_path = str(path/f'{args.model_arch}.onnx')
input_names = [ "actual_input_1" ] + [ "learned_%d" % i for i in range(16) ]
output_names = [ "output1" ]
torch.onnx.export(learn.model, dummy_input, output_path, verbose=True, input_names=input_names, output_names=output_names)
```
-------
Also as a side question, how should I expect performance to compare with Neo compilation starting with a model saved via ONNX versus JIT?