diff --git a/OnnxStack.Core/Config/ExecutionProvider.cs b/OnnxStack.Core/Config/ExecutionProvider.cs index beccce8c..e108feac 100644 --- a/OnnxStack.Core/Config/ExecutionProvider.cs +++ b/OnnxStack.Core/Config/ExecutionProvider.cs @@ -6,6 +6,7 @@ public enum ExecutionProvider Cuda = 1, Cpu = 2, CoreML = 3, - OpenVino = 4 + OpenVino = 4, + TensorRT = 5 } } diff --git a/OnnxStack.Core/Extensions/Extensions.cs b/OnnxStack.Core/Extensions/Extensions.cs index 3cd55b9d..6359bc3c 100644 --- a/OnnxStack.Core/Extensions/Extensions.cs +++ b/OnnxStack.Core/Extensions/Extensions.cs @@ -51,6 +51,9 @@ public static SessionOptions GetSessionOptions(this OnnxModelConfig configuratio }; sessionOptions.AppendExecutionProvider_OpenVINO(deviceId); return sessionOptions; + case ExecutionProvider.TensorRT: + sessionOptions.AppendExecutionProvider_Tensorrt(configuration.DeviceId.Value); + return sessionOptions; } }