From 1fb0aac97f097f9fe0c1ace5869655106bb0937e Mon Sep 17 00:00:00 2001 From: David Corvoysier Date: Wed, 24 Jul 2024 14:13:59 +0000 Subject: [PATCH] fix(qbits): correct stride when packing --- optimum/quanto/tensor/qbits/qbits.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/optimum/quanto/tensor/qbits/qbits.py b/optimum/quanto/tensor/qbits/qbits.py index 2f1f82b8..bf1049e1 100644 --- a/optimum/quanto/tensor/qbits/qbits.py +++ b/optimum/quanto/tensor/qbits/qbits.py @@ -165,9 +165,14 @@ def dequantize(self): @staticmethod def load_from_state_dict(state_dict, prefix, qtype, axis, group_size, size, stride): - data_size = size if group_size is None else grouped_shape(size, axis, group_size) - # In row major, inner dimension (stride 1) is the last one - data_stride = (data_size[1], 1) + if group_size is None: + data_size = size + data_stride = stride + else: + data_size = grouped_shape(size, axis, group_size) + assert len(data_size) == 2 + # In row major, inner dimension (stride 1) is the last one + data_stride = (data_size[1], 1) inner_tensors_dict = { "_data": PackedTensor.load_from_state_dict( state_dict, prefix + "_data.", qtype.bits, data_size, data_stride