From ff2306fac46865ea029793c761df63fb43c820c8 Mon Sep 17 00:00:00 2001
From: menglcai <menglcai@amd.com>
Date: Mon, 16 Dec 2024 16:45:11 +0800
Subject: [PATCH] update allocator when execute operator

---
 .../dml/DmlExecutionProvider/src/ExecutionProvider.cpp      | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp
index 6b0faaad43175..618321a789fa9 100644
--- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp
+++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp
@@ -444,6 +444,12 @@ namespace Dml
     {
         ORT_TRY
         {
+
+        if (m_allocator)
+        {
+            m_context->SetAllocator(m_allocator);
+        }
+
         assert(!m_closed);
 
         DML_BINDING_DESC persistentResourceBindingDesc =