update int4 usage

Signed-off-by: Kaihui-intel <[email protected]>
intel · Jul 22, 2024 · e538a63 · WeiweiZhang1 · Jul 23, 2024 · e538a63
1 parent 967b780
commit e538a63
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/neural_compressor/torch/quantization/algorithm_entry.py b/neural_compressor/torch/quantization/algorithm_entry.py
@@ -567,9 +567,14 @@ def autoround_quantize_entry(
         if quant_config.name != AUTOROUND or quant_config.dtype == "fp32":
             continue
         else:
+            dtype = quant_config.dtype
+            bits = quant_config.bits
+            if dtype != "int" and "int" in dtype:
+                bits = int(dtype.lstrip("int"))
+                dtype = "int"
             weight_config[op_name] = {
-                "data_type": quant_config.dtype,
-                "bits": quant_config.bits,
+                "data_type": dtype,
+                "bits": bits,
                 "sym": quant_config.use_sym,
                 "group_size": quant_config.group_size,
                 "act_bits": quant_config.act_bits,