Skip to content

Commit

Permalink
Fix pylint errors
Browse files Browse the repository at this point in the history
Signed-off-by: Seokjun An <[email protected]>
  • Loading branch information
quic-seokan committed Jan 23, 2025
1 parent e2669b7 commit 74831ba
Showing 1 changed file with 4 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
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)
Expand Down Expand Up @@ -204,6 +205,9 @@ def wrapper(*args, **kwargs):


def export(model: torch.nn.Module, *args, **kwargs):
"""
Export a torch model to ONNX with precomputed scale and offset.
"""
if not isinstance(model, torch.nn.Module):
raise NotImplementedError

Expand All @@ -215,8 +219,6 @@ def export(model: torch.nn.Module, *args, **kwargs):

@contextmanager
def _precompute_encodings(model: torch.nn.Module):
from aimet_torch.quantization.base import QuantizerBase

with ExitStack() as stack:
for q in model.modules():
if isinstance(q, QuantizerBase):
Expand Down

0 comments on commit 74831ba

Please sign in to comment.