Fix WeightOnlyLinear bits type when dtype="intx" (intel#1827)

yiliu30 · May 30, 2024 · 19ff13e · 19ff13e
1 parent 5dafe5f
commit 19ff13e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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