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 77b4782e5..b29c83809 100644
--- a/src/brevitas_examples/llm/llm_quant/prepare_for_quantize.py
+++ b/src/brevitas_examples/llm/llm_quant/prepare_for_quantize.py
@@ -24,8 +24,7 @@ def _set_bert_mha_attributes(module):
 
 
 _SET_ATTRIBUTES_MAP = {
-    BertAttention: _set_bert_mha_attributes,
-}
+    BertAttention: _set_bert_mha_attributes,}
 
 
 def set_mha_attributes(model):