fix quant proxy

Xilinx · Aug 28, 2024 · fe38f1c · fe38f1c
1 parent be81c4e
commit fe38f1c
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/src/brevitas/proxy/parameter_quant.py b/src/brevitas/proxy/parameter_quant.py
@@ -125,8 +125,6 @@ def forward(self, x: torch.Tensor) -> Union[Tensor, QuantTensor]:
                 if not self.training and self.cache_inference_quant_weight and self._cached_weight is None:
                     self._cached_weight = self.cache_class(
                         out.detach(), metadata_only=self.cache_inference_quant_weight_metadata_only)
-                else:
-                    out = out[0]
         else:  # quantization disabled
             out = x
         return out