Update src/brevitas_examples/llm/llm_quant/learned_round_utils.py

Co-authored-by: Pablo Monteagudo Lago <[email protected]>
Xilinx · Dec 3, 2024 · a15163b · a15163b
1 parent acee4d3
commit a15163b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/brevitas_examples/llm/llm_quant/learned_round_utils.py b/src/brevitas_examples/llm/llm_quant/learned_round_utils.py
@@ -124,7 +124,7 @@ def apply_learned_round(
     lr_scheduler_kwargs: Optional[Dict] = None,
     learned_round_loss_kwargs: Optional[Dict] = None,
     scale_optimizer_class: Optional[str] = None,
-    scale_optimizer_kwargs: Optional[dict] = None,
+    scale_optimizer_kwargs: Optional[Dict] = None,
 ) -> None:
     # Parse strings to obtain the arguments for the optimizer
     learned_round = parse_learned_round(learned_round)