How to save a model to TRT file? Not covered in the docs

[talcs]

I'd like to save my model into TRT file. Could you please provide a sample code?

[jaybdub]

Hi @talcs,

Thanks for reaching out.

Assuming you built your model like

model_trt = torch2trt(...)

You can then save the raw TensorRT engine like this.

with open("model.plan", 'wb') as f:
    f.write(model_trt.engine.serialize())

The input and output names will be

print(model_trt.input_names)
print(model_trt.output_names)

Let me know if this helps or you run into any issues.

Best,
John