Fix

Xilinx · Oct 17, 2023 · bb8f0ad · bb8f0ad
1 parent 200c0de
commit bb8f0ad
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/brevitas/utils/quant_utils.py b/src/brevitas/utils/quant_utils.py
@@ -5,11 +5,12 @@
 from brevitas.core.function_wrapper import *
 from brevitas.core.quant import RescalingIntQuant
 from brevitas.inject.enum import FloatToIntImplType
+from brevitas.proxy.quant_proxy import QuantProxyFromInjector
 
 
 def _is_proxy_in_export_mode(model):
     for submodule in model.modules():
-        if hasattr(submodule, 'export_mode'):
+        if isinstance(submodule, QuantProxyFromInjector) and hasattr(submodule, 'export_mode'):
             return submodule.export_mode
     return False