diff --git a/src/brevitas_examples/llm/llm_quant/bias_corr.py b/src/brevitas_examples/llm/llm_quant/bias_corr.py
index 900777874..dc603f8a3 100644
--- a/src/brevitas_examples/llm/llm_quant/bias_corr.py
+++ b/src/brevitas_examples/llm/llm_quant/bias_corr.py
@@ -10,6 +10,6 @@
 
 @torch.no_grad()
 def apply_bias_correction(model, dataloader):
-    with bias_correction_mode(curr_layer):
+    with bias_correction_mode(model):
         for inps in dataloader:
             model(**inps)