Repository: pytorch
Issue: [export] Torch custom class export issue
Labels: oncall: pt2,oncall: export

### 🐛 Describe the bug

We are trying to implement the functionality to export custom classes within Torch-TensorRT. TensorRT engines are expressed as custom classes (https://github.com/pytorch/TensorRT/blob/re-export/core/runtime/register_jit_hooks.cpp#L77-L83)
Following https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA/edit
1) I implement the flatten call here : https://github.com/pytorch/TensorRT/blob/re-export/core/runtime/TRTEngine.cpp#L321-L336
2) I register a fake class with the unflatten call : https://github.com/pytorch/TensorRT/blob/re-export/py/torch_tensorrt/dynamo/runtime/register_fake_class.py#L16-L21

When I re-export the graph with these compiled TRT engines, I face an error as shown in this [reexport.log]
(https://github.com/user-attachments/files/17438975/reexport.log) 

Error msg: 
```py
Traceback (most recent call last):
  File "/home/dperi/Downloads/TensorRT/test.py", line 29, in <module>
    torch_tensorrt.save(trt_gm, "./trt.ep", inputs=inputs, retrace=True)
  File "/home/dperi/Downloads/TensorRT/py/torch_tensorrt/_compile.py", line 534, in save
    exp_program = torch.export.export(
                  ^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/export/__init__.py", line 366, in export
    return _export(
           ^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/export/_trace.py", line 1013, in wrapper
    raise e
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/export/_trace.py", line 986, in wrapper
    ep = fn(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/export/exported_program.py", line 116, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/export/_trace.py", line 1954, in _export
    export_artifact = export_func(  # type: ignore[operator]
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/export/_trace.py", line 1744, in _non_strict_export
    aten_export_artifact = _to_aten_func(  # type: ignore[operator]
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/export/_trace.py", line 642, in _export_to_aten_ir
    gm, graph_signature = transform(aot_export_module)(
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/export/_trace.py", line 1674, in _aot_export_non_strict
    gm, sig = aot_export(wrapped_mod, args, kwargs=kwargs, **flags)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 1262, in aot_export_module
    fx_g, metadata, in_spec, out_spec = _aot_export_function(
                                        ^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 1497, in _aot_export_function
    fx_g, meta = create_aot_dispatcher_function(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 524, in create_aot_dispatcher_function
    return _create_aot_dispatcher_function(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 625, in _create_aot_dispatcher_function
    fw_metadata = run_functionalized_fw_and_collect_metadata(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/collect_metadata_analysis.py", line 194, in inner
    flat_f_outs = f(*flat_f_args)
                  ^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/utils.py", line 184, in flat_fn
    tree_out = fn(*args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 863, in functional_call
    out = mod(*args[params_len:], **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/export/_trace.py", line 1657, in forward
    tree_out = torch.fx.Interpreter(self._export_root).run(
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/fx/interpreter.py", line 146, in run
    self.env[node] = self.run_node(node)
                     ^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/fx/interpreter.py", line 203, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/fx/interpreter.py", line 320, in call_module
    return submod(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/Downloads/TensorRT/py/torch_tensorrt/_features.py", line 56, in wrapper
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/home/dperi/Downloads/TensorRT/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py", line 275, in forward
    outputs: List[torch.Tensor] = torch.ops.tensorrt.execute_engine(
                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/_ops.py", line 1122, in __call__
    return _call_overload_packet_from_python(self, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/_ops.py", line 1134, in _call_overload_packet_from_python
    torch_function_called, ret = torch._C._maybe_call_torch_function_for_op_packet(
                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/_export/non_strict_utils.py", line 551, in __torch_function__
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/_ops.py", line 1122, in __call__
    return _call_overload_packet_from_python(self, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/.pyenv/versions/3.11.7/lib/python3.11/site-packages/torch/_ops.py", line 1167, in _call_overload_packet_from_python
    raise RuntimeError(err_msg)
RuntimeError: Fail to match any TorchBindOverload of tensorrt.execute_engine with following exceptions:
Overload name default:
 tensorrt::execute_engine() Expected a value of type '__torch__.torch.classes.tensorrt.Engine (of Python compilation unit at: 0)' for argument '_1' but instead found type 'FakeScriptObject'.
Position: 1
Value: Torch-TensorRT TensorRT Engine:
  Name: _run_on_acc_0_engine
  Inputs: [
    id: 0
      name: x
      shape: [1, 10]
      dtype: Float
  ]
  Outputs: [
    id: 0
      name: output0
      shape: [1, 40]
      dtype: Float
  ]
  Device: Device(ID: 0, Name: NVIDIA GeForce RTX 3080 Ti, SM Capability: 8.6, Type: GPU)
  Hardware Compatibility: Disabled
  Target Platform: linux_x86_64

Declaration: tensorrt::execute_engine(Tensor[] _0, __torch__.torch.classes.tensorrt.Engine _1) -> Tensor[] _0
Cast error details: _1 is expected to be a FakeScriptObject of __torch__.torch.classes.tensorrt.Engine
```

During the re-export, this forward function gets called https://github.com/pytorch/TensorRT/blob/re-export/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py#L275 and the self.engine object here is a FakeScriptObject 

```py
(Pdb) self.engine
<torch._library.fake_class_registry.FakeScriptObject object at 0x7557289bf010>
(Pdb) self.engine.real_obj
Torch-TensorRT TensorRT Engine:
  Name: _run_on_acc_0_engine
  Inputs: [
    id: 0
      name: x
      shape: [1, 10]
      dtype: Float
  ]
  Outputs: [
    id: 0
      name: output0
      shape: [1, 40]
      dtype: Float
  ]
  Device: Device(ID: 0, Name: NVIDIA GeForce RTX 3080 Ti, SM Capability: 8.6, Type: GPU)
  Hardware Compatibility: Disabled
  Target Platform: linux_x86_64

(Pdb) input_tensors
[FunctionalTensor(_to_functional_tensor(FakeTensor(..., device='cuda:0', size=(1, 10)),
       device='cuda:0'))]
```
Any suggestions on how to resolve the issue in the log ? 
cc: @angelayi 

### Versions

[pip3] pytorch-lightning==2.0.7
[pip3] pytorch_sphinx_theme==0.0.24
[pip3] pytorch-triton==3.1.0+cf34004b8a
[pip3] pytorch-triton-rocm==2.2.0
[pip3] pytorchvideo==0.1.5
[pip3] torch==2.6.0.dev20241011+cu124
[pip3] torch_tensorrt==2.6.0.dev0+cb03ca14a
[pip3] torchmetrics==1.4.0.post0
[pip3] torchprofile==0.0.4
[pip3] torchsurgeon==0.1.2
[pip3] torchvision==0.20.0.dev20241011+cu124
[pip3] triton==3.0.0

cc @ezyang @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4