better check?

Xilinx · Aug 16, 2024 · 0f4225b · 0f4225b
1 parent ac98fc2
commit 0f4225b
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/brevitas/proxy/quant_proxy.py b/src/brevitas/proxy/quant_proxy.py
@@ -85,9 +85,10 @@ def __init__(self, quant_layer: nn.Module, quant_injector: Injector) -> None:
         self.tracked_module_list = []
         self.add_tracked_module(quant_layer)
         self.disable_quant = False
-        self.is_zero_zero_point = type(
-            self.quant_injector.zero_point_impl
-        ) is ZeroZeroPoint if 'zero_point_impl' in self.quant_injector else False
+        self.is_zero_zero_point = False
+        for m in self.modules():
+            if isinstance(m, ZeroZeroPoint):
+                self.is_zero_zero_point = True
         self.is_signed = self.quant_injector.signed if 'signed' in quant_injector else None
 
     @property