提交 3348e5a7 编写于 作者: L lianliguang

deal something special of adam's kernel select

上级 ce9f7734
......@@ -82,6 +82,13 @@ bool IsValidKernelInfo(const std::shared_ptr<CNode> &kernel_node, const kernel::
}
return true;
};
if (AnfAlgo::GetCNodeName(kernel_node) == "Adam") {
auto input_num = AnfAlgo::GetInputTensorNum(kernel_node);
if (AnfAlgo::GetPrevNodeOutputFormat(kernel_node, input_num - 1) !=
kernel_build_info.GetInputFormat(input_num - 1)) {
return false;
}
}
if (AnfAlgo::GetCNodeName(kernel_node) == prim::kPrimCast->name()) {
return AnfAlgo::GetOutputInferDataType(kernel_node, 0) == kernel_build_info.GetOutputDeviceType(0) &&
AnfAlgo::GetPrevNodeOutputInferDataType(kernel_node, 0) == kernel_build_info.GetInputDeviceType(0);
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册