fix bug under sparse mode

mindspore-lab · Sep 25, 2023 · 6400436 · 6400436
1 parent a46885c
commit 6400436
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/mindspore_rec/ops/embedding.py b/mindspore_rec/ops/embedding.py
@@ -194,7 +194,10 @@ def construct(self, indices):
             weight_flatten = self.gather_revert(weight_unique, unique_idx, 0)
             out = self.reshape(weight_flatten, shape)
         else:
-            out = self.embedding_table.get(indices)
+            shape = self.shape(indices) + (self.embedding_size,)
+            indices_flatten = self.reshape_first(indices, (-1,))
+            out = self.map_tensor_get(self.embedding_table, indices_flatten)
+            out = self.reshape(out, shape)
 
         if self.max_norm is not None:
             axis = _make_axis_range(F.rank(indices), F.rank(out))