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