Update vqlinear.py (#103)

fix index unpack size
microsoft · Oct 31, 2024 · ace1393 · ace1393
1 parent fde03de
commit ace1393
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/vptq/layers/vqlinear.py b/vptq/layers/vqlinear.py
@@ -532,9 +532,9 @@ def dequant(self):
             indices, res_indices = self.unpack_index_tensor(
                 pack_tensor=self.indices,
                 index_bits=index_bits,
-                num_elements=self.in_features,
+                num_elements=self.group_size,
                 res_bits=index_res_bits,
-                num_res_elements=self.in_features,
+                num_res_elements=self.group_size,
                 index_dtype=torch.uint16,
             )