From d628376c1d7c04c924a993f69da21c5a8152ac9b Mon Sep 17 00:00:00 2001 From: Nick Fraser Date: Thu, 23 May 2024 10:59:30 +0100 Subject: [PATCH] Fix (llm): Add all rewriters to the list --- src/brevitas_examples/llm/llm_quant/prepare_for_quantize.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/brevitas_examples/llm/llm_quant/prepare_for_quantize.py b/src/brevitas_examples/llm/llm_quant/prepare_for_quantize.py index e96a6d946..2a9505227 100644 --- a/src/brevitas_examples/llm/llm_quant/prepare_for_quantize.py +++ b/src/brevitas_examples/llm/llm_quant/prepare_for_quantize.py @@ -13,7 +13,7 @@ def replace_mha_with_quantizable_layers(model, dtype): for src_module, (quantizable_module, quantizable_module_kwargs) in QUANTIZABLE_MHA_MAP.items(): rewriter = ModuleToModuleByClass( src_module, quantizable_module, **quantizable_module_kwargs, dtype=dtype) - rewriters.append(rewriter) + rewriters.append(rewriter) if not rewriters: warnings.warn( f"No module to replace was found. Supported modules are {list(QUANTIZABLE_MHA_MAP.keys())}"