fix get outputaddress error

This commit is contained in:
kswang 2021-06-24 10:43:53 +08:00
parent 4932854776
commit 6afd18a4d8
1 changed files with 3 additions and 0 deletions

View File

@ -1598,6 +1598,9 @@ void SessionBasic::UpdateOutputTensors(const VectorRef *outputs,
if (iter != tensor_to_node.end()) {
const auto &node = iter->second.first;
const auto &output_index = iter->second.second;
if (!AnfAlgo::OutputAddrExist(node, output_index)) {
continue;
}
const auto &address = AnfAlgo::GetMutableOutputAddr(node, output_index);
tensor->set_device_address(address);