Function torch_tensorrt::torchscript::embed_engine_in_new_module ¶
-
Defined in File torch_tensorrt.h
Function Documentation ¶
-
TORCHTRT_API torch::jit::Module
torch_tensorrt::torchscript
::
embed_engine_in_new_module
( const std::string & engine , Device device ) ¶
-
Take a previously created TensorRT engine and embed it in in a TorchScript module.
Takes a pre-built serialized TensorRT engine and embeds it in a TorchScript module. Registers execution of the engine as the forward method of the module Forward is defined as: forward(Tensor[]) -> Tensor[]
- Parameters
-
-
engine
: std::string - Pre-built serialized TensorRT engine -
device
: CompileSepc::Device - Device information
-
TensorRT bindings must have names with the following format:
-
[symbol].[index in input / output array] ex.
-
[x.0, x.1, x.2] -> [y.0]
- Return
-
: A new module targeting a TensorRT engine