Col_lin
/
mindspore2022
forked from huawei/mindspore2022
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

adapt input to attr between cpu and aicpu embeddinglookup

This commit is contained in:
wuxuejian 2020-08-11 15:49:42 +08:00
parent f92735d14a
commit 5e6b1d42da
2 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
mindspore/ccsrc/backend/optimizer/pass/convert_const_input_to_attr.cc
View File

@ -50,6 +50,12 @@ const AnfNodePtr ConvertConstInputToAttr::Process(const FuncGraphPtr &, const An
if (!ConstInputToAttrInfoRegistry::Instance().GetRegisterByOpName(AnfAlgo::GetCNodeName(cnode), &reg)) {
continue;
}
if (AnfAlgo::GetCNodeName(cnode) == prim::kPrimEmbeddingLookup->name() ||
AnfAlgo::GetCNodeName(cnode) == prim::kPrimEmbeddingLookupCommGrad->name()) {
if (!AnfAlgo::HasNodeAttr(kAttrPrimitiveTarget, cnode)) {
continue;
}
}
ConstInputToAttr(cnode, reg.GetConstInputAttrInfo());
}
return node;

2
tests/st/ops/cpu/test_embedding_look_up_op.py
View File

@ -26,7 +26,7 @@ context.set_context(mode=context.GRAPH_MODE, device_target="CPU")
class Net(nn.Cell):
def __init__(self, offset):
super(Net, self).__init__()
self.embedding = P.EmbeddingLookup()
self.embedding = P.EmbeddingLookup().add_prim_attr("primitive_target", "CPU")
self.offset = offset
def construct(self, param, index):