From 023d60aee45e79bd7f5fa038c0f1d56ca94440d5 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Thu, 10 Sep 2020 15:14:50 +0800 Subject: [PATCH] fix(mgb/atlas): use aclmdlGetOutputSizeByIndex to create DataBuffer GitOrigin-RevId: df7f1b225dd4e39443e4ebf94339f6f47de55102 --- src/opr/impl/atlas_runtime_op.cpp | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/opr/impl/atlas_runtime_op.cpp b/src/opr/impl/atlas_runtime_op.cpp index ddc645e6c..e2fb4ff3a 100644 --- a/src/opr/impl/atlas_runtime_op.cpp +++ b/src/opr/impl/atlas_runtime_op.cpp @@ -308,6 +308,8 @@ void AtlasRuntimeOpr::scn_do_execute() { for (size_t i = 0; i < input().size(); i++) { auto value_pair = input_getter.get(batch, i); auto input_size = aclmdlGetInputSizeByIndex(m_model_desc, i); + //! FIXME iff enable dynamic batchsize and dynamic aipp, the input + //! size should be the size of aclmdlGetInputSizeByIndex. if (enable_dynamic_batch) { mgb_assert(input_size == value_pair.second / batch * m_dyn_batch_choices[0], @@ -345,8 +347,12 @@ void AtlasRuntimeOpr::scn_do_execute() { "failed to create atlas output dataset."); for (size_t i = 0; i < nr_outputs; i++) { auto value_pair = output_getter.get(batch, i); + size_t output_size = value_pair.second; + if (enable_dynamic_batch) { + output_size = aclmdlGetOutputSizeByIndex(m_model_desc, i); + } aclDataBuffer* output_db = - aclCreateDataBuffer(value_pair.first, value_pair.second); + aclCreateDataBuffer(value_pair.first, output_size); mgb_assert(output_db != nullptr, "failed to create atlas output data buffer for output " "%zu:%s.", -- GitLab