add sharded + device_map check

Author: a-r-r-o-w

src/diffusers/models/modeling_utils.py

@@ -803,6 +803,12 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                     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."
+                        )
+
                     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