diff options
Diffstat (limited to 'src/backends/backendsCommon/TensorHandleFactoryRegistry.cpp')
-rw-r--r-- | src/backends/backendsCommon/TensorHandleFactoryRegistry.cpp | 14 |
1 files changed, 14 insertions, 0 deletions
diff --git a/src/backends/backendsCommon/TensorHandleFactoryRegistry.cpp b/src/backends/backendsCommon/TensorHandleFactoryRegistry.cpp index 0670461b54..cc8a1361a3 100644 --- a/src/backends/backendsCommon/TensorHandleFactoryRegistry.cpp +++ b/src/backends/backendsCommon/TensorHandleFactoryRegistry.cpp @@ -49,6 +49,20 @@ ITensorHandleFactory* TensorHandleFactoryRegistry::GetFactory(ITensorHandleFacto return nullptr; } +ITensorHandleFactory* TensorHandleFactoryRegistry::GetFactory(ITensorHandleFactory::FactoryId id, + MemorySource memSource) const +{ + for (auto& factory : m_Factories) + { + if (factory->GetId() == id && factory->GetImportFlags() == static_cast<MemorySourceFlags>(memSource)) + { + return factory.get(); + } + } + + return nullptr; +} + void TensorHandleFactoryRegistry::AquireMemory() { for (auto& mgr : m_MemoryManagers) |