From 5c6910ed9cc0880b1dede1a128d921bdaba2b1a2 Mon Sep 17 00:00:00 2001
From: Ryan Hill <ryanhill@microsoft.com>
Date: Thu, 22 Apr 2021 23:07:56 -0700
Subject: [PATCH] Fix memory cleanup on unload

---
 onnxruntime/core/providers/cuda/cuda_execution_provider.h | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.h b/onnxruntime/core/providers/cuda/cuda_execution_provider.h
index 827c675fd8..095da844da 100644
--- a/onnxruntime/core/providers/cuda/cuda_execution_provider.h
+++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.h
@@ -172,7 +172,11 @@ class CUDAExecutionProvider : public IExecutionProvider {
 
   struct ContextCacheHolder {
     ContextCacheHolder() {
-      RunOnUnload([this] { p.reset(); });
+      // Keep a weak pointer to the object, if the weak pointer can be locked, then the shared pointer is still around, so we can reset it
+      RunOnUnload([&, weak_p_ = std::weak_ptr<PerThreadContextMap>(p)] {
+        if (auto lock = weak_p_.lock())
+          p.reset();
+      });
     }
     std::shared_ptr<PerThreadContextMap> p{std::make_shared<PerThreadContextMap>()};
   };