def test_device_map(self):
        """
        Test if the quantized model int4 weight-only is working properly with "auto" and custom device maps.
        The custom device map performs cpu/disk offloading as well. Also verifies that the device map is
        correctly set (in the `hf_device_map` attribute of the model).
        """
        custom_device_map_dict = {
            "time_text_embed": torch_device,
            "context_embedder": torch_device,
            "x_embedder": torch_device,
            "transformer_blocks.0": "cpu",
            "single_transformer_blocks.0": "disk",
            "norm_out": torch_device,
            "proj_out": "cpu",
        }
        device_maps = ["auto", custom_device_map_dict]

        inputs = self.get_dummy_tensor_inputs(torch_device)
        # requires with different expected slices since models are different due to offload (we don't quantize modules offloaded to cpu/disk)
        expected_slice_auto = np.array(
            [
                0.34179688,
                -0.03613281,
                0.01428223,
                -0.22949219,
                -0.49609375,
                0.4375,
                -0.1640625,
                -0.66015625,
                0.43164062,
            ]
        )
        expected_slice_offload = np.array(
            [0.34375, -0.03515625, 0.0123291, -0.22753906, -0.49414062, 0.4375, -0.16308594, -0.66015625, 0.43554688]
        )
        for device_map in device_maps:
            if device_map == "auto":
                expected_slice = expected_slice_auto
            else:
                expected_slice = expected_slice_offload
            with tempfile.TemporaryDirectory() as offload_folder:
                quantization_config = TorchAoConfig("int4_weight_only", group_size=64)
                quantized_model = FluxTransformer2DModel.from_pretrained(
                    "hf-internal-testing/tiny-flux-pipe",
                    subfolder="transformer",
                    quantization_config=quantization_config,
                    device_map=device_map,
                    torch_dtype=torch.bfloat16,
                    offload_folder=offload_folder,
                )

                weight = quantized_model.transformer_blocks[0].ff.net[2].weight

                # Note that when performing cpu/disk offload, the offloaded weights are not quantized, only the weights on the gpu.
                # This is not the case when the model are already quantized
                if "transformer_blocks.0" in device_map:
                    self.assertTrue(isinstance(weight, nn.Parameter))
                else:
                    self.assertTrue(isinstance(weight, AffineQuantizedTensor))

                output = quantized_model(**inputs)[0]
                output_slice = output.flatten()[-9:].detach().float().cpu().numpy()
                self.assertTrue(numpy_cosine_similarity_distance(output_slice, expected_slice) < 2e-3)

            with tempfile.TemporaryDirectory() as offload_folder:
                quantization_config = TorchAoConfig("int4_weight_only", group_size=64)
                quantized_model = FluxTransformer2DModel.from_pretrained(
                    "hf-internal-testing/tiny-flux-sharded",
                    subfolder="transformer",
                    quantization_config=quantization_config,
                    device_map=device_map,
                    torch_dtype=torch.bfloat16,
                    offload_folder=offload_folder,
                )

                weight = quantized_model.transformer_blocks[0].ff.net[2].weight
                if "transformer_blocks.0" in device_map:
                    self.assertTrue(isinstance(weight, nn.Parameter))
                else:
                    self.assertTrue(isinstance(weight, AffineQuantizedTensor))

                output = quantized_model(**inputs)[0]
                output_slice = output.flatten()[-9:].detach().float().cpu().numpy()
                self.assertTrue(numpy_cosine_similarity_distance(output_slice, expected_slice) < 2e-3)

        Test if the quantized model int4 weight-only is working properly with "auto" and custom device maps.
        The custom device map performs cpu/disk offloading as well. Also verifies that the device map is
        correctly set (in the `hf_device_map` attribute of the model).
        
tests/quantization/torchao/test_torchao.py