From 19ff13e8a8963744349e46013ef522fcb3e8c3d8 Mon Sep 17 00:00:00 2001 From: "Wang, Chang" Date: Thu, 30 May 2024 16:49:00 +0800 Subject: [PATCH] Fix WeightOnlyLinear bits type when dtype="intx" (#1827) --- neural_compressor/torch/algorithms/weight_only/modules.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/neural_compressor/torch/algorithms/weight_only/modules.py b/neural_compressor/torch/algorithms/weight_only/modules.py index 1eaab289c08..768aa0c2fdc 100644 --- a/neural_compressor/torch/algorithms/weight_only/modules.py +++ b/neural_compressor/torch/algorithms/weight_only/modules.py @@ -66,7 +66,7 @@ def __init__( self.use_optimum_format = use_optimum_format self.dtype = dtype if self.dtype != "int" and "int" in self.dtype: # for nf4, fp4 - bits = self.dtype.lstrip("int") + bits = int(self.dtype.lstrip("int")) self.dtype = "int" if "int" not in self.dtype: # for nf4, fp4 from neural_compressor.torch.algorithms.weight_only.utility import FLOAT_MAPPING, INT_MAPPING