refine details

zehao-intel · zehao-intel · commit aa54c5f268b5 · 2024-05-15T16:00:33.000+08:00
Signed-off-by: zehao-intel &lt;zehao.huang@intel.com&gt;
diff --git a/neural_compressor/torch/quantization/algorithm_entry.py b/neural_compressor/torch/quantization/algorithm_entry.py
@@ -502,7 +502,7 @@ def fp8_quant_entry(
         return model
 
 
-###################### FP16 Algo Entry ##################################
+###################### Mixed Precision Algo Entry ##################################
 @register_algo(MIX_PRECISION)
 def mix_precision_entry(
     model: torch.nn.Module, configs_mapping: Dict[Tuple[str], MixPrecisionConfig], *args, **kwargs
diff --git a/neural_compressor/torch/quantization/config.py b/neural_compressor/torch/quantization/config.py
@@ -1173,7 +1173,7 @@ def register_supported_configs(cls) -> List[OperatorConfig]:
         supported_configs = []
         mix_precision_config = MixPrecisionConfig(
             dtype=["fp16", "fp32"],
-            device=["cpu", "cuda"],
+            device=["auto", "cpu", "cuda"],
         )
         operators = cls.supported_fp16_ops
         supported_configs.append(OperatorConfig(config=mix_precision_config, operators=operators))

Original file line number	Diff line number	Diff line change
`@@ -1173,7 +1173,7 @@ def register_supported_configs(cls) -> List[OperatorConfig]:`
`1173`	`1173`	`supported_configs = []`
`1174`	`1174`	`mix_precision_config = MixPrecisionConfig(`
`1175`	`1175`	`dtype=["fp16", "fp32"],`
`1176`		`- device=["cpu", "cuda"],`
	`1176`	`+ device=["auto", "cpu", "cuda"],`
`1177`	`1177`	`)`
`1178`	`1178`	`operators = cls.supported_fp16_ops`
`1179`	`1179`	`supported_configs.append(OperatorConfig(config=mix_precision_config, operators=operators))`