From 93e3c6abfc9f5abfc72d7f8c76d7dfa4b3fb6867 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Wed, 3 Jul 2024 10:04:59 -0700 Subject: [PATCH] Fix an import after previous refactor Summary: att Test Plan: CI Reviewers: Subscribers: Tasks: Tags: --- torchao/dtypes/affine_quantized_tensor.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/torchao/dtypes/affine_quantized_tensor.py b/torchao/dtypes/affine_quantized_tensor.py index 666f7c1947..d4f607a8f4 100644 --- a/torchao/dtypes/affine_quantized_tensor.py +++ b/torchao/dtypes/affine_quantized_tensor.py @@ -509,9 +509,10 @@ def __torch_dispatch__(cls, func, types, args, kwargs): def get_plain(self): from torchao.quantization.quant_primitives import ( ZeroPointDomain, - unpack_tinygemm_scales_and_zeros, quantize_affine, ) + from torchao.quantization.utils import unpack_tinygemm_scales_and_zeros + cur_shape = self.shape assert len(cur_shape) == 4 inner_k_tiles = cur_shape[-1] * 2