Browse Source

!8755 fix embeddinglookup omission

From: @yao_yf
Reviewed-by: @kisnwang,@stsuteng
Signed-off-by: @stsuteng
tags/v1.1.0
mindspore-ci-bot Gitee 5 years ago
parent
commit
ed288f5151
1 changed files with 1 additions and 1 deletions
  1. +1
    -1
      mindspore/nn/layer/embedding.py

+ 1
- 1
mindspore/nn/layer/embedding.py View File

@@ -253,7 +253,7 @@ class EmbeddingLookup(Cell):
shp = self.shape(indices) + (self.embedding_size,) shp = self.shape(indices) + (self.embedding_size,)
indices_flatten = self.reshape(indices, (-1,)) indices_flatten = self.reshape(indices, (-1,))
unique_id, unique_idx = self.unique(indices_flatten) unique_id, unique_idx = self.unique(indices_flatten)
weight_unique = self.gatherv2(unique_id)
weight_unique = self.gatherv2(self.embedding_table, unique_id, 0)
weight_flatten = self.gather_revert(weight_unique, unique_idx, 0) weight_flatten = self.gather_revert(weight_unique, unique_idx, 0)
out = self.reshape(weight_flatten, shp) out = self.reshape(weight_flatten, shp)
else: else:


Loading…
Cancel
Save