diff --git a/winml/lib/Api/LearningModelSession.cpp b/winml/lib/Api/LearningModelSession.cpp index c4f5862f43..9aab61d47c 100644 --- a/winml/lib/Api/LearningModelSession.cpp +++ b/winml/lib/Api/LearningModelSession.cpp @@ -126,7 +126,8 @@ void LearningModelSession::Initialize() { // Register the custom operator registry auto model = model_.as(); - WINML_THROW_IF_FAILED(session->RegisterCustomRegistry(model->GetOperatorRegistry())); + operatorRegistry_.reset(model->GetOperatorRegistry()); + WINML_THROW_IF_FAILED(session->RegisterCustomRegistry(operatorRegistry_.get())); // Register only the transformers not already in ORT session->RegisterGraphTransformers(); diff --git a/winml/lib/Api/LearningModelSession.h b/winml/lib/Api/LearningModelSession.h index 6d4b2af80a..5e2bc84761 100644 --- a/winml/lib/Api/LearningModelSession.h +++ b/winml/lib/Api/LearningModelSession.h @@ -6,6 +6,7 @@ #include "LearningModelSession.g.h" #include "LearningModelBinding.h" +#include "MLOperatorAuthor.h" #include "WinML_Lock.h" #include "WinMLAdapter.h" @@ -102,6 +103,12 @@ struct LearningModelSession : LearningModelSessionT { private: com_ptr inference_session_; + struct IMLOperatorRegistryDeleter { + void operator()(IMLOperatorRegistry* p) { + p->Release(); + } + }; + std::unique_ptr operatorRegistry_; // reference to the active execution provider. weak onnxruntime::IExecutionProvider* cached_execution_provider_ = nullptr;