Fix init quant (#136)

89f3feae · Casper · GitHub · eea08aa6 · 89f3feae
Unverified Commit 89f3feae authored Nov 02, 2023 by Casper Committed by GitHub Nov 02, 2023
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

awq/quantize/quantizer.py awq/quantize/quantizer.py +1 -1

No files found.
--- a/awq/quantize/quantizer.py
+++ b/awq/quantize/quantizer.py
@@ -332,7 +332,7 @@ class AwqQuantizer:
        
        clear_memory()
        
-        if "attention_mask" in layer_kwargs.keys():
+        if layer_kwargs.get("attention_mask") is not None:
            layer_kwargs["attention_mask"] = layer_kwargs["attention_mask"].to("cuda")

        return modules, layer_kwargs, inps