1 files changed, 2 insertions, 2 deletions
diff --git a/candle-transformers/src/generation/mod.rs b/candle-transformers/src/generation/mod.rs
index d95a0595..85ffb59c 100644
--- a/candle-transformers/src/generation/mod.rs
+++ b/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::{DType, Error, Result, Tensor};
+use candle::{Context, DType, Error, Result, Tensor};
 use rand::{distributions::Distribution, SeedableRng};
 
 #[derive(Clone, PartialEq, Debug)]
@@ -45,7 +45,7 @@ impl LogitsProcessor {
             .enumerate()
             .max_by(|(_, u), (_, v)| u.total_cmp(v))
             .map(|(i, _)| i as u32)
-            .unwrap();
+            .context("empty logits")?;
         Ok(next_token)
     }