From e0d1cf19a6e54867dc5f15a2977d3685fbdb417c Mon Sep 17 00:00:00 2001
From: Ryan Hill <38674843+RyanUnderhill@users.noreply.github.com>
Date: Wed, 2 Sep 2020 21:21:18 -0700
Subject: [PATCH] Fix allocator bug (#5042)

---
 onnxruntime/core/framework/provider_bridge_ort.cc | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/onnxruntime/core/framework/provider_bridge_ort.cc b/onnxruntime/core/framework/provider_bridge_ort.cc
index 0fe2a7cba1..0ed4288368 100644
--- a/onnxruntime/core/framework/provider_bridge_ort.cc
+++ b/onnxruntime/core/framework/provider_bridge_ort.cc
@@ -187,6 +187,13 @@ struct Provider_IExecutionProvider_Router_Impl : Provider_IExecutionProvider_Rou
     return std::make_shared<Provider_IAllocator_Impl>(IExecutionProvider::GetAllocator(id, mem_type));
   }
 
+  AllocatorPtr GetAllocator(int id, OrtMemType mem_type) const override {
+    auto allocator = outer_->Provider_GetAllocator(id, mem_type);
+    if (!allocator)
+      return nullptr;
+    return static_cast<Provider_IAllocator_Impl*>(allocator.get())->p_;
+  }
+
   std::unique_ptr<Provider_IDataTransfer> Provider_GetDataTransfer() const override {
     return std::unique_ptr<Provider_IDataTransfer>(reinterpret_cast<Provider_IDataTransfer*>(IExecutionProvider::GetDataTransfer().release()));
   }