fix: support ITensor amex in quantize and add export_torch_mode() wrapper

[zewenli98]

Description

1. support ITensor amex in quantize
2. add export_torch_mode() wrapper in torch_tensorrt.compile to modelopt quant models

Fixes #4429

Type of change

Please delete options that are not relevant and/or add your own.

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[narendasan]

Can you explain what this context is and maybe add some tests?

[zewenli98]

Can you explain what this context is and maybe add some tests?

1. In quantize(), amex could be trt.ITensor which doesn't have amax.numel()
2. ModelOpt requires with export_torch_mode() wrapper to make inserted quantizers traceable while exporting pytorch models. See: torch_tensorrt.compile() fails with "fake tensor in exported program constant's list" for ModelOpt quantized models unless wrapped in export_torch_mode() #4429

[narendasan]

Can you explain what this context is and maybe add some tests?

1. In quantize(), `amex` could be trt.ITensor which doesn't have `amax.numel()`

2. ModelOpt requires `with export_torch_mode()` wrapper to make inserted quantizers traceable while exporting pytorch models. See: [torch_tensorrt.compile() fails with "fake tensor in exported program constant's list" for ModelOpt quantized models unless wrapped in export_torch_mode() #4429](https://github.com/pytorch/TensorRT/issues/4429)

I see, is there any harm in always having this context? should we catch cases where the modelopt is missing but the context is needed and throw an error/warning?

[zewenli98]

I see, is there any harm in always having this context? should we catch cases where the modelopt is missing but the context is needed and throw an error/warning?

I remember ModelOpt is not required in torch-trt. If we need to check all models with ModelOpt API, the package should be required rather than opt in.