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

Giuseppe5 · May 31, 2024 · 5ceb7e8 · 5ceb7e8
1 parent 32f9ffe
commit 5ceb7e8
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())}"