Change condition

Signed-off-by: quic-sanising <[email protected]>

Author: quic-sanising

QEfficient/transformers/sampler/sampler.py

@@ -218,7 +218,7 @@ def sampler_forward(
 
     logits = logits.reshape(-1, vocab_size)  # Reshape tensor to 2D
 
-    if input_ids.shape[1] != spec_length:  # Prefill phase, initialize retained states
+    if input_ids.shape[1] > spec_length:  # Prefill phase, initialize retained states
         repetition_penalty_retain_state_selected = torch.mul(repetition_penalty_retain_state_selected, 0)
         presence_penalty_retain_state_selected = torch.mul(presence_penalty_retain_state_selected, 0)
         # TODO: Replace scatter_ with CtxScatterFunc; Replace -1 with int_max while exporting on onnx