Use index-select for the embeddings as it supports backprop. (#298)

1 files changed, 1 insertions, 1 deletions

diff --git a/candle-nn/src/embedding.rs b/candle-nn/src/embedding.rs
index 050123be..f4ba88e7 100644
--- a/candle-nn/src/embedding.rs
+++ b/candle-nn/src/embedding.rs
@@ -23,7 +23,7 @@ impl Embedding {
         let mut final_dims = indexes.dims().to_vec();
         final_dims.push(self.hidden_size);
         let indexes = indexes.flatten_all()?;
-        let values = Tensor::embedding(&indexes, &self.embeddings)?;
+        let values = self.embeddings.index_select(&indexes, 0)?;
         let values = values.reshape(final_dims)?;
         Ok(values)
     }