summaryrefslogtreecommitdiff
path: root/candle-examples/examples/bert/main.rs
diff options
context:
space:
mode:
authorLaurent Mazare <laurent.mazare@gmail.com>2023-07-22 11:39:27 +0200
committerGitHub <noreply@github.com>2023-07-22 10:39:27 +0100
commit43c72232927ca80c850a73ce977c2063d5a2dcf5 (patch)
treec93c07984e06b1925313f4f641a8b1a3956fc0ed /candle-examples/examples/bert/main.rs
parent52c5d8c087f6a2ee91b807467860eb3e96bb6267 (diff)
downloadcandle-43c72232927ca80c850a73ce977c2063d5a2dcf5.tar.gz
candle-43c72232927ca80c850a73ce977c2063d5a2dcf5.tar.bz2
candle-43c72232927ca80c850a73ce977c2063d5a2dcf5.zip
Rename the .r functions to .dims so as to be a bit more explicit. (#220)
Diffstat (limited to 'candle-examples/examples/bert/main.rs')
-rw-r--r--candle-examples/examples/bert/main.rs2
1 files changed, 1 insertions, 1 deletions
diff --git a/candle-examples/examples/bert/main.rs b/candle-examples/examples/bert/main.rs
index 33f0a1fe..6672ad09 100644
--- a/candle-examples/examples/bert/main.rs
+++ b/candle-examples/examples/bert/main.rs
@@ -161,7 +161,7 @@ fn main() -> Result<()> {
let embeddings = model.forward(&token_ids, &token_type_ids)?;
println!("generated embeddings {:?}", embeddings.shape());
// Apply some avg-pooling by taking the mean embedding value for all tokens (including padding)
- let (_n_sentence, n_tokens, _hidden_size) = embeddings.shape().r3()?;
+ let (_n_sentence, n_tokens, _hidden_size) = embeddings.dims3()?;
let embeddings = (embeddings.sum(1)? / (n_tokens as f64))?;
println!("pooled embeddings {:?}", embeddings.shape());
let mut similarities = vec![];