Fix

Xilinx · Jul 30, 2024 · 69e5e98 · 69e5e98
1 parent 718288f
commit 69e5e98
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/src/brevitas/proxy/groupwise_float_runtime_quant.py b/src/brevitas/proxy/groupwise_float_runtime_quant.py
@@ -4,6 +4,7 @@
 
 from brevitas.proxy.float_runtime_quant import ActFloatQuantProxyFromInjectorBase
 from brevitas.quant_tensor import GroupwiseFloatQuantTensor
+from brevitas.quant_tensor import QuantTensor
 from brevitas.utils.quant_utils import _CachedIOGroupwiseFloat
 
 
@@ -17,13 +18,11 @@ def group_dim(self):
     def group_size(self):
         return self.quant_injector.group_size
 
-    def forward(
-            self, x: Union[Tensor,
-                           GroupwiseFloatQuantTensor]) -> Union[Tensor, GroupwiseFloatQuantTensor]:
+    def forward(self, x: Union[Tensor, QuantTensor]) -> Union[Tensor, GroupwiseFloatQuantTensor]:
         out = x
         if self.fused_activation_quant_proxy is not None:
             y = x
-            if isinstance(y, GroupwiseFloatQuantTensor):
+            if isinstance(y, QuantTensor):
                 y = y.value
 
             if self.export_mode: