From b9d8e77131250545c640f4f3e3210d6c4c876797 Mon Sep 17 00:00:00 2001 From: liuqi Date: Fri, 9 Mar 2018 15:32:35 +0800 Subject: [PATCH] Fix multiple outputs' memory reuse logic bug. --- mace/core/workspace.cc | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/mace/core/workspace.cc b/mace/core/workspace.cc index 1e927c3e..2cb5e237 100644 --- a/mace/core/workspace.cc +++ b/mace/core/workspace.cc @@ -137,16 +137,17 @@ void Workspace::CreateImageOutputTensor(const NetDef &net_def) { for (auto &op : net_def.op()) { if (!op.mem_id().empty()) { auto mem_ids = op.mem_id(); - for (auto mem_id : mem_ids) { + int count = mem_ids.size(); + for (int i = 0; i < count; ++i) { std::unique_ptr tensor - (new Tensor(preallocated_allocator_.GetBuffer(mem_id), dtype)); + (new Tensor(preallocated_allocator_.GetBuffer(mem_ids[i]), dtype)); tensor->SetSourceOpName(op.name()); VLOG(3) << "Tensor: " << op.name() << "(" << op.type() << ")" << "; Mem: " - << mem_id << "; Image shape: " + << mem_ids[i] << "; Image shape: " << dynamic_cast(tensor->UnderlyingBuffer())->image_shape()[0] << ", " << dynamic_cast(tensor->UnderlyingBuffer())->image_shape()[1]; - tensor_map_[op.output(0)] = std::move(tensor); + tensor_map_[op.output(i)] = std::move(tensor); } } } -- GitLab