Review

src/brevitas_examples/llm/main.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import argparse
+from contextlib import nullcontext
 from copy import deepcopy
 import functools
 import sys
@@ -425,9 +426,15 @@ def quantize_llm(args):
                 new_funct = functools.partial(update_internal_dict, m)
                 m._hf_hook.post_forward = hooked_on_a_function(m._hf_hook.post_forward, new_funct)
 
-    with functional_quantization_mode(
-            model, {torch.nn.functional.scaled_dot_product_attention: ScaledDotProductAttention},
-            enabled=args.functional_sdpa_quant):
+    # If we are doing functional SDPA quantization, we create the correct context manager,
+    # otherwise nullcontext. We would love to avoid the extra indentation level but it doesn't seem easy.
+    if args.functional_sdpa_quant:
+        quantization_cm = functional_quantization_mode(
+            model, {torch.nn.functional.scaled_dot_product_attention: ScaledDotProductAttention})
+    else:
+        quantization_cm = nullcontext()
+
+    with quantization_cm:
         with torch.no_grad():
             model(**calibration_loader[0])