Update

[ghstack-poisoned]
pytorch · Feb 5, 2025 · 26850da · 26850da
1 parent aaba2d8
commit 26850da
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/torchao/quantization/quant_api.py b/torchao/quantization/quant_api.py
@@ -779,6 +779,7 @@ class Int4WeightOnlyConfig(AOBaseConfig):
     use_hqq: bool = False
     zero_point_domain: Optional[ZeroPointDomain] = ZeroPointDomain.NONE
 
+
 # for BC
 # TODO maybe change other callsites
 int4_weight_only = Int4WeightOnlyConfig
@@ -812,7 +813,9 @@ def _int4_weight_only_transform(
     quant_max = 15
     eps = 1e-6
     preserve_zero = LAYOUT_TO_PRESERVE_ZEROS[type(layout)]
-    zero_point_dtype = torch.bfloat16
+    zero_point_dtype = (
+        weight.dtype if isinstance(layout, Int4CPULayout) else torch.bfloat16
+    )
 
     # nonlocal zero_point_domain
     assert (