Merge pull request #88 from casper-hansen/fix_padding_mask

Fix unexpected keyword

Merge pull request #88 from casper-hansen/fix_padding_mask
Fix unexpected keyword
a4ea423d · Casper · GitHub · 3fa74007 · d016c513 · a4ea423d
Unverified Commit a4ea423d authored Oct 02, 2023 by Casper Committed by GitHub Oct 02, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

awq/modules/fused/attn.py awq/modules/fused/attn.py +2 -1

No files found.
--- a/awq/modules/fused/attn.py
+++ b/awq/modules/fused/attn.py
@@ -146,7 +146,8 @@ class QuantAttentionFused(nn.Module):
    def forward(
        self,
-        hidden_states:torch.Tensor, past_key_value=None, attention_mask=None, position_ids=None, output_attentions=False, use_cache=False
+        hidden_states:torch.Tensor, past_key_value=None, attention_mask=None, position_ids=None, 
+        output_attentions=False, use_cache=False, *args, **kwargs
    ):
        bsz, seqlen, _ = hidden_states.shape
        if bsz != self.cache_batch_size: