forked from mindspore-Ecosystem/mindspore
set value node & parameter's device dtype to the node connected with's device info
This commit is contained in:
parent
aef9c4d838
commit
934958290e
|
@ -176,7 +176,7 @@ void SetTensorDeviceInfo(const kernel::KernelBuildInfo &selected_kernel_info, co
|
|||
if (AnfAlgo::GetOutputDeviceDataType(real_input_node, 0) == kTypeUnknown || is_ref) {
|
||||
std::vector<std::string> output_format = {selected_kernel_info.GetInputFormat(input_index)};
|
||||
builder->SetOutputsFormat(output_format);
|
||||
std::vector<TypeId> output_type = {AnfAlgo::GetOutputInferDataType(real_input_node, 0)};
|
||||
std::vector<TypeId> output_type = {AnfAlgo::GetInputDeviceDataType(kernel_node, input_index)};
|
||||
builder->SetOutputsDeviceType(output_type);
|
||||
AnfAlgo::SetSelectKernelBuildInfo(builder->Build(), real_input_node.get());
|
||||
}
|
||||
|
|
|
@ -298,7 +298,10 @@ CNodePtr InsertCastForInput(const FuncGraphPtr &func_graph, const CNodePtr &cnod
|
|||
auto cur_input = AnfAlgo::GetInputNode(cnode, input_index);
|
||||
auto kernel_with_index = AnfAlgo::VisitKernel(cur_input, 0);
|
||||
auto is_weight_boundary = [](const AnfNodePtr &node) -> bool {
|
||||
if (node->isa<ValueNode>() || node->isa<Parameter>()) {
|
||||
if (node->isa<ValueNode>()) {
|
||||
return true;
|
||||
}
|
||||
if (node->isa<Parameter>() && AnfAlgo::IsParameterWeight(node->cast<ParameterPtr>())) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
|
|
Loading…
Reference in New Issue