From aba9ba4ef0250765ccf6d61c8b20f7d8d022f5a9 Mon Sep 17 00:00:00 2001 From: limingqi107 Date: Thu, 21 Oct 2021 16:58:22 +0800 Subject: [PATCH] fix the weight update bug in the shared weight scene --- mindspore/ccsrc/runtime/framework/actor/data_prepare_actor.cc | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mindspore/ccsrc/runtime/framework/actor/data_prepare_actor.cc b/mindspore/ccsrc/runtime/framework/actor/data_prepare_actor.cc index 680c0d13a36..ac9e6a32b01 100644 --- a/mindspore/ccsrc/runtime/framework/actor/data_prepare_actor.cc +++ b/mindspore/ccsrc/runtime/framework/actor/data_prepare_actor.cc @@ -421,7 +421,6 @@ void DataPrepareActor::PrepareDataForWeightNode(const AnfNodePtr &backend_node, UpdateRefCount(host_tensor_address.get(), true); } MS_EXCEPTION_IF_NULL(host_tensor_address); - DeviceTensorStore::GetInstance().Insert(front_node.get(), host_tensor_address); if (host_tensor_address->DeviceType() == device_tensor->DeviceType()) { AnfAlgo::SetOutputAddr(host_tensor_address, 0, backend_node.get()); } else { @@ -429,6 +428,9 @@ void DataPrepareActor::PrepareDataForWeightNode(const AnfNodePtr &backend_node, << ", device tensor type:" << device_tensor->DeviceType(); } } + // Maybe the same host_tensor_address corresponds to the different front_node in shared weight scene, + // so need update the device tensor store always. + DeviceTensorStore::GetInstance().Insert(front_node.get(), host_tensor_address); // If the ptr of device tensor is not nullptr, it indicates that the device data has been prepared. MS_EXCEPTION_IF_NULL(host_tensor_address);