Fix (example/llm): Add all rewriters to the list (#956)

Xilinx · May 23, 2024 · b535615 · b535615
1 parent eeffd2b
commit b535615
Showing 1 changed file with 1 addition and 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
@@ -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())}"