Fix (quant_tensor): Produce valid IntQuantTensor after AvgPool functi…

…onal call
Xilinx · Oct 4, 2024 · af37fcc · af37fcc
1 parent f84a0e1
commit af37fcc
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/brevitas/quant_tensor/int_torch_handler.py b/src/brevitas/quant_tensor/int_torch_handler.py
@@ -114,6 +114,7 @@ def avg_pool2d_handler(
         avg_scaling = kernel_size * kernel_size
 
     quant_input = quant_input.set(value=x)
+    quant_input = quant_input.set(scale=quant_input.scale / avg_scaling)
     quant_input = quant_input.set(bit_width=max_acc_bit_width(quant_input.bit_width, avg_scaling))
     return quant_input
 
@@ -133,6 +134,7 @@ def adaptive_avg_pool2d_handler(quant_input, output_shape):
     reduce_size = reduce(mul, k_size, 1)
 
     quant_input = quant_input.set(value=x)
+    quant_input = quant_input.set(scale=quant_input.scale / reduce_size)
     quant_input = quant_input.set(bit_width=max_acc_bit_width(quant_input.bit_width, reduce_size))
     return quant_input