From 236d6c6de416866eaa574ded10affc7a28a601e4 Mon Sep 17 00:00:00 2001 From: kswang Date: Tue, 16 Jun 2020 17:13:03 +0800 Subject: [PATCH] fix cpu reshape bug --- mindspore/ccsrc/device/cpu/cpu_kernel_runtime.cc | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/mindspore/ccsrc/device/cpu/cpu_kernel_runtime.cc b/mindspore/ccsrc/device/cpu/cpu_kernel_runtime.cc index c8bfd9f7cd..63b4749047 100644 --- a/mindspore/ccsrc/device/cpu/cpu_kernel_runtime.cc +++ b/mindspore/ccsrc/device/cpu/cpu_kernel_runtime.cc @@ -161,8 +161,12 @@ BaseRef CPUKernelRuntime::CreatTensorForOutput(const AnfNodePtr &input_node, siz } tensor::TensorPtr tensor = std::make_shared(type_id, temp_shape); MS_EXCEPTION_IF_NULL(tensor); - address->ptr_ = tensor->data_c(true); - address->ref_count_ = INIT_NODE_REF; + if (address->ref_count_ > 0 && address->ptr_ != nullptr) { + tensor->set_device_address(address); + } else { + address->ptr_ = tensor->data_c(true); + address->ref_count_ = INIT_NODE_REF; + } tensor->set_dirty(false); return tensor; } else if (input_node->isa() || input_node->isa()) { @@ -211,6 +215,7 @@ void CPUKernelRuntime::BindInputOutput(const session::KernelGraph *kernel_graph, } tensor->set_dirty(true); } + address->ref_count_ = INIT_NODE_REF; tensor->set_device_address(address); } @@ -220,7 +225,7 @@ void CPUKernelRuntime::BindInputOutput(const session::KernelGraph *kernel_graph, // new output and bind ptr auto output_nodes = kernel_graph->outputs(); for (const auto &item : output_nodes) { - auto item_with_index = AnfAlgo::VisitKernelWithReturnType(item, 0); + auto item_with_index = AnfAlgo::VisitKernelWithReturnType(item, 0, true); auto out = CreatTensorForOutput(item_with_index.first, item_with_index.second, input_map); outputs->push_back(std::move(out)); }