apply review suggestions

Author: a-r-r-o-w

src/diffusers/models/modeling_utils.py

@@ -718,11 +718,9 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
             hf_quantizer = None
 
         if hf_quantizer is not None:
-            is_bnb_quantization_method = hf_quantizer.quantization_config.quant_method.value == "bitsandbytes"
-            is_torchao_quantization_method = hf_quantizer.quantization_config.quant_method.value == "torchao"
-            if (is_bnb_quantization_method or is_torchao_quantization_method) and device_map is not None:
+            if device_map is not None:
                 raise NotImplementedError(
-                    "Currently, `device_map` is automatically inferred for quantized bitsandbytes and torchao models. Support for providing `device_map` as an input will be added in the future."
+                    "Currently, `device_map` is automatically inferred for quantized models. Support for providing `device_map` as an input will be added in the future."
                 )
 
             hf_quantizer.validate_environment(torch_dtype=torch_dtype, from_flax=from_flax, device_map=device_map)

tests/quantization/torchao/test_torchao.py

@@ -418,7 +418,7 @@ def test_torch_compile(self):
             quantization_config = TorchAoConfig("int8_weight_only")
             components = self.get_dummy_components(quantization_config, model_id=model_id)
             pipe = FluxPipeline(**components)
-            pipe.to(device=torch_device, dtype=torch.bfloat16)
+            pipe.to(device=torch_device)
 
             inputs = self.get_dummy_inputs(torch_device)
             normal_output = pipe(**inputs)[0].flatten()[-32:]