From 6172a60eaa5fadad9cbd475522db579aade4085a Mon Sep 17 00:00:00 2001 From: lichun Date: Thu, 13 May 2021 10:20:42 +0800 Subject: [PATCH] fix tbe handle not unregisted --- tests/ut/ge/hybrid/ge_hybrid_unittest.cc | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/ut/ge/hybrid/ge_hybrid_unittest.cc b/tests/ut/ge/hybrid/ge_hybrid_unittest.cc index 7cd8a30a..a27222ef 100644 --- a/tests/ut/ge/hybrid/ge_hybrid_unittest.cc +++ b/tests/ut/ge/hybrid/ge_hybrid_unittest.cc @@ -76,6 +76,10 @@ static ge::OpDescPtr CreateOpDesc(string name = "", string type = "") { } TEST_F(UtestGeHybrid, aicore_op_task_init_success) { + ComputeGraphPtr graph = std::make_shared("test"); + GeRootModelPtr ge_root_model = make_shared(graph); + HybridModel hybrid_model(ge_root_model); + // build aicore task auto aicore_task = std::unique_ptr(new(std::nothrow)hybrid::AiCoreOpTask()); domi::TaskDef task_def; @@ -99,7 +103,7 @@ TEST_F(UtestGeHybrid, aicore_op_task_init_success) { op_desc->SetExtAttr(ge::OP_EXTATTR_NAME_TBE_KERNEL, tbe_kernel); std::string kernel_name("kernel/Add"); AttrUtils::SetStr(op_desc, op_desc->GetName() + "_kernelname", kernel_name); - ASSERT_EQ(aicore_task->InitWithTaskDef(*op_desc.get(), task_def), SUCCESS); + ASSERT_EQ(aicore_task->InitWithTaskDef(hybrid_model, *op_desc.get(), task_def), SUCCESS); rtStream_t stream = nullptr; rtStreamCreate(&stream, 0); ASSERT_EQ(aicore_task->LaunchKernel(stream), SUCCESS);