From b2f4e1f7722b7d9a9c5948f4b7d3e4cb32c1525b Mon Sep 17 00:00:00 2001 From: "Carson M." Date: Sun, 7 Jul 2024 12:26:55 -0500 Subject: [PATCH] fix: create an environment if one does not exist when registering TensorRT EP. ref #236 --- src/execution_providers/tensorrt.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/execution_providers/tensorrt.rs b/src/execution_providers/tensorrt.rs index 1ea8dd8..599eaf0 100644 --- a/src/execution_providers/tensorrt.rs +++ b/src/execution_providers/tensorrt.rs @@ -230,6 +230,11 @@ impl ExecutionProvider for TensorRTExecutionProvider { fn register(&self, session_builder: &SessionBuilder) -> Result<()> { #[cfg(any(feature = "load-dynamic", feature = "tensorrt"))] { + // The TensorRT execution provider specifically is pretty picky about requiring an environment to be initialized by the + // time we register it. This isn't always the case in `ort`, so if we get to this point, let's make sure we have an + // environment initialized. + let _ = crate::get_environment(); + let mut trt_options: *mut ort_sys::OrtTensorRTProviderOptionsV2 = std::ptr::null_mut(); crate::error::status_to_result(crate::ortsys![unsafe CreateTensorRTProviderOptions(&mut trt_options)]).map_err(Error::ExecutionProvider)?; let (key_ptrs, value_ptrs, len, keys, values) = super::map_keys! {