diff --git a/candle-transformers/src/utils.rs b/candle-transformers/src/utils.rs
index 50d3b707..3cbcac5c 100644
--- a/candle-transformers/src/utils.rs
+++ b/candle-transformers/src/utils.rs
@@ -3,9 +3,13 @@ use candle::{Result, Tensor};
 pub fn apply_repeat_penalty(logits: &Tensor, penalty: f32, context: &[u32]) -> Result<Tensor> {
     let device = logits.device();
     let mut logits = logits.to_vec1::<f32>()?;
-    let context: std::collections::HashSet<_> = context.iter().collect();
-    for (token_id, logit) in logits.iter_mut().enumerate() {
-        if context.contains(&(token_id as u32)) {
+    let mut already_seen = std::collections::HashSet::new();
+    for token_id in context {
+        if already_seen.contains(token_id) {
+            continue;
+        }
+        already_seen.insert(token_id);
+        if let Some(logit) = logits.get_mut(*token_id as usize) {
             if *logit >= 0. {
                 *logit /= penalty
             } else {