fix quantization_config init

Signed-off-by: changwangss <[email protected]>
changwangss · Nov 29, 2024 · 0b36916 · 0b36916
1 parent 7e73e82
commit 0b36916
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/optimum/intel/neural_compressor/quantization.py b/optimum/intel/neural_compressor/quantization.py
@@ -392,11 +392,14 @@ def _weight_only_quantization(
                 )
                 low_cpu_mem_usage = False
                 model = model_class.from_pretrained(model_id, low_cpu_mem_usage=low_cpu_mem_usage, **loading_kwargs)
-            quantization_config.update(**{"device": "xpu"})
-            quantization_config.post_init_xpu()
         else:
             model = model_class.from_pretrained(model_id, low_cpu_mem_usage=low_cpu_mem_usage, **loading_kwargs)
-            quantization_config.post_init_cpu()
+
+    if use_xpu:
+        quantization_config.update(**{"device": "xpu"})
+        quantization_config.post_init_xpu()
+    else:
+        quantization_config.post_init_cpu()
 
     model.config.update({"low_cpu_mem_usage": low_cpu_mem_usage})
     model.eval()