@ptrblck Thank you for help I tried the export approach but with torch_tensorrt.dynamo.export (torch.export does not support TorchTRT) but it I got the following exception with that approach:
[ERROR] Preprocessing failed: Detected that you are using FX to torch.jit.trace a dynamo-optimized function. This is not supported at the moment.