Skip to content

Commit

Permalink
fix: create an environment if one does not exist when registering Ten…
Browse files Browse the repository at this point in the history
…sorRT EP. ref #236
  • Loading branch information
decahedron1 committed Jul 7, 2024
1 parent bb57252 commit b2f4e1f
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions src/execution_providers/tensorrt.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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! {
Expand Down

0 comments on commit b2f4e1f

Please sign in to comment.