From bb8f0ad47ba1813ff51b2f09e99d33f1979a7b06 Mon Sep 17 00:00:00 2001 From: Giuseppe Franco Date: Tue, 17 Oct 2023 15:56:18 +0100 Subject: [PATCH] Fix --- src/brevitas/utils/quant_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/brevitas/utils/quant_utils.py b/src/brevitas/utils/quant_utils.py index f63bda16c..9c09bc176 100644 --- 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