Use Tensor::argmax instead of manual cpu impl (#3173)

Author: ivarflakstad

candle-transformers/src/generation/mod.rs

@@ -3,7 +3,7 @@
 //! Functionality for modeling sampling strategies and logits processing in text generation
 //! with support for temperature-based sampling, top-k filtering, nucleus sampling (top-p),
 //! and combinations thereof.
-use candle::{Context, DType, Error, Result, Tensor};
+use candle::{DType, Error, Result, Tensor};
 use rand::{distr::Distribution, SeedableRng};
 
 #[derive(Clone, PartialEq, Debug)]
@@ -41,19 +41,12 @@ impl LogitsProcessor {
     }
 
     fn sample_argmax(&mut self, logits: Tensor) -> Result<u32> {
-        let logits_v: Vec<f32> = logits.to_vec1()?;
-        let next_token = logits_v
-            .iter()
-            .enumerate()
-            .max_by(|(_, u), (_, v)| u.total_cmp(v))
-            .map(|(i, _)| i as u32)
-            .context("empty logits")?;
-        Ok(next_token)
+        logits.argmax(candle::D::Minus1)?.to_scalar::<u32>()
     }
 
     fn sample_gumbel_softmax(&mut self, logits: &Tensor, temperature: f64) -> Result<u32> {
         let sampled = candle_nn::sampling::gumbel_softmax(logits, temperature, candle::D::Minus1)?;
-        sampled.to_vec0::<u32>()
+        sampled.to_scalar::<u32>()
     }
 
     fn sample_multinomial(&mut self, prs: &Vec<f32>) -> Result<u32> {