fix

Author: a-r-r-o-w

src/diffusers/models/modeling_utils.py

@@ -802,13 +802,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                     revision=revision,
                     subfolder=subfolder or "",
                 )
-                if hf_quantizer is not None:
-                    is_torchao_quantization_method = quantization_config.quant_method == QuantizationMethod.TORCHAO
-                    if device_map is not None and is_torchao_quantization_method:
-                        raise NotImplementedError(
-                            "Loading sharded checkpoints, while passing `device_map`, is not supported with `torchao` quantization. This will be supported in the near future."
-                        )
-
+                if hf_quantizer is not None and is_bnb_quantization_method:
                     model_file = _merge_sharded_checkpoints(sharded_ckpt_cached_folder, sharded_metadata)
                     logger.info("Merged sharded checkpoints as `hf_quantizer` is not None.")
                     is_sharded = False