Minor fixes

pytorch · Nov 14, 2024 · b9245cc · b9245cc
1 parent 6ff85f5
commit b9245cc
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 4 deletions.
diff --git a/test/hqq/test_hqq_affine.py b/test/hqq/test_hqq_affine.py
@@ -1,6 +1,6 @@
 import unittest
 import torch
-from torchao.quantization.quant_primitives import (
+from torchao.quantization import (
     ZeroPointDomain,
     MappingType,
 )

diff --git a/torchao/dtypes/__init__.py b/torchao/dtypes/__init__.py
@@ -1,6 +1,5 @@
 from . import affine_quantized_tensor_ops
 
-# from ..prototype.dtypes.uint2 import UInt2Tensor, BitnetTensor
 from .affine_quantized_tensor import (
     AffineQuantizedTensor,
     MarlinQQQTensor,
@@ -29,7 +28,6 @@
     PlainLayout,
 )
 
-# from ..prototype.dtypes.uint2 import UInt2Tensor, BitnetTensor
 
 __all__ = [
     "NF4Tensor",

diff --git a/torchao/prototype/hqq/example.py b/torchao/prototype/hqq/example.py
@@ -3,7 +3,7 @@
 from torchao.dtypes.affine_quantized_tensor import (
     to_affine_quantized_intx,
 )
-from torchao.quantization.quant_primitives import (
+from torchao.quantization import (
         ZeroPointDomain,
         MappingType,
 )