From 13c2eb3524b77700e2700442fefe830c18c6b08e Mon Sep 17 00:00:00 2001
From: Tiago Koji Castro Shibata <tiago.shibata@gmail.com>
Date: Mon, 16 Dec 2019 15:41:54 -0800
Subject: [PATCH] Fix leaking operator registry (#2645)

Fix https://microsoft.visualstudio.com/OS/_workitems/edit/24354916
---
 winml/lib/Api/LearningModelSession.cpp | 3 ++-
 winml/lib/Api/LearningModelSession.h   | 7 +++++++
 2 files changed, 9 insertions(+), 1 deletion(-)

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<winmlp::LearningModel>();
-  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<LearningModelSession> {
 
  private:
   com_ptr<winmla::IInferenceSession> inference_session_;
+  struct IMLOperatorRegistryDeleter {
+    void operator()(IMLOperatorRegistry* p) {
+        p->Release();
+    }
+  };
+  std::unique_ptr<IMLOperatorRegistry, IMLOperatorRegistryDeleter> operatorRegistry_;
 
   // reference to the active execution provider. weak
   onnxruntime::IExecutionProvider* cached_execution_provider_ = nullptr;