Fix unittest failure

Signed-off-by: Seokjun An <[email protected]>
quic · Jan 23, 2025 · e223fb4 · e223fb4
1 parent 74831ba
commit e223fb4
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/v2/experimental/onnx/_export.py b/TrainingExtensions/torch/src/python/aimet_torch/v2/experimental/onnx/_export.py
@@ -46,7 +46,6 @@
 from torch.onnx import is_in_onnx_export, symbolic_helper
 
 from aimet_torch.v2.utils import patch_attr
-from aimet_torch.quantization.base import QuantizerBase
 
 
 aimet_opset = onnxscript.values.Opset(domain="aimet", version=1)
@@ -219,6 +218,8 @@ def export(model: torch.nn.Module, *args, **kwargs):
 
 @contextmanager
 def _precompute_encodings(model: torch.nn.Module):
+    # pylint: disable=import-outside-toplevel
+    from aimet_torch.quantization.base import QuantizerBase
     with ExitStack() as stack:
         for q in model.modules():
             if isinstance(q, QuantizerBase):