Fix position ids (#215)

4d49ae1f · Casper · GitHub · 8dbf009b · 4d49ae1f
Unverified Commit 4d49ae1f authored Dec 03, 2023 by Casper Committed by GitHub Dec 03, 2023
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

awq/utils/fused_utils.py awq/utils/fused_utils.py +1 -1

No files found.
--- a/awq/utils/fused_utils.py
+++ b/awq/utils/fused_utils.py
@@ -32,7 +32,7 @@ def prepare_input_ids(input_ids: torch.Tensor, last_forward_num_tokens: int):
        num_new_tokens = num_input_tokens - last_forward_num_tokens
        
        # after context is processed, slice to latest token
-        if num_new_tokens in [0,1]:
+        if num_new_tokens == 1:
            input_ids = input_ids[:, -1:]

    return input_ids, last_forward_num_tokens + num_new_tokens