Fix (nn/TruncAvgPool): Remove any quant tensor manual manipulation.

Author: nickfraser

src/brevitas/nn/quant_avg_pool.py

@@ -58,6 +58,7 @@ def _avg_scaling(self):
         else:
             return self.kernel_size * self.kernel_size
 
+    # TODO: Replace with functional call
     def forward(self, input: Union[Tensor, QuantTensor]):
         x = self.unpack_input(input)
 
@@ -71,8 +72,6 @@ def forward(self, input: Union[Tensor, QuantTensor]):
             if not isinstance(x, QuantTensor):
                 x = self.cache_class.quant_tensor.set(value=x)
             y = AvgPool2d.forward(self, x)
-            rescaled_value = y.value * self._avg_scaling
-            y = y.set(value=rescaled_value)
             y = self.trunc_quant(y)
         else:
             y = AvgPool2d.forward(self, _unpack_quant_tensor(x))
@@ -123,6 +122,7 @@ def compute_kernel_size_stride(input_shape, output_shape):
             stride_list.append(stride)
         return kernel_size_list, stride_list
 
+    # TODO: Replace with functional call
     def forward(self, input: Union[Tensor, QuantTensor]):
         x = self.unpack_input(input)
 
@@ -139,10 +139,6 @@ def forward(self, input: Union[Tensor, QuantTensor]):
             if not isinstance(x, QuantTensor):
                 x = self.cache_class.quant_tensor.set(value=x)
             y = AdaptiveAvgPool2d.forward(self, x)
-            k_size, stride = self.compute_kernel_size_stride(x.value.shape[2:], y.value.shape[2:])
-            reduce_size = reduce(mul, k_size, 1)
-            rescaled_value = y.value * reduce_size  # remove avg scaling
-            y = y.set(value=rescaled_value)
             y = self.trunc_quant(y)
         else:
             y = AdaptiveAvgPool2d.forward(self, _unpack_quant_tensor(x))