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[]
TensorRT bindings must have names with the following format:
-
[symbol].[index in input / output array] ex.
-
[x.0, x.1, x.2] -> [y.0]
- Parameters
-
-
engine – std::string - Pre-built serialized TensorRT engine
-
device – CompileSepc::Device - Device information
-
- Returns
-
: A new module targeting a TensorRT engine
-