Correct weight_quant names

Xilinx · Jul 7, 2023 · 2eb5b56 · 2eb5b56
1 parent 425f945
commit 2eb5b56
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/brevitas_examples/imagenet_classification/ptq/ptq_common.py b/src/brevitas_examples/imagenet_classification/ptq/ptq_common.py
@@ -242,18 +242,19 @@ def act_kwargs_prefix(prefix, kwargs):
         if act_quant_type == 'asym':
             act_quant = act_quant.let(**{'low_percentile_q': 100 - act_quant_percentile})
 
-    weight_quant_and_bit_width = {'quant': weight_quant, 'bit_width': weight_bit_width}
+    weight_quant_and_bit_width = {
+        'weight_quant': weight_quant, 'weight_bit_width': weight_bit_width}
     act_quant_and_bit_width = {'act_quant': act_quant, 'bit_width': act_bit_width}
     sym_act_quant_and_bit_width = {'act_quant': sym_act_quant, 'bit_width': act_bit_width}
     per_tensor_act_quant_and_bit_width = {
         'act_quant': per_tensor_act_quant, 'bit_width': act_bit_width}
     quant_linear_kwargs = {
-        **weight_kwargs_prefix('weight_', weight_quant_and_bit_width),
+        **weight_quant_and_bit_width,
         'dtype': dtype,
         'return_quant_tensor': weight_return_quant_tensor,
         'bias_quant': bias_quant}
     quant_conv_kwargs = {
-        **weight_kwargs_prefix('weight_', weight_quant_and_bit_width),
+        **weight_quant_and_bit_width,
         'dtype': dtype,
         'return_quant_tensor': weight_return_quant_tensor,
         'bias_quant': bias_quant}