def test_torch_compile(self):
        r"""Test that verifies if torch.compile works with torchao quantization."""
        for model_id in ["hf-internal-testing/tiny-flux-pipe", "hf-internal-testing/tiny-flux-sharded"]:
            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)

            inputs = self.get_dummy_inputs(torch_device)
            normal_output = pipe(**inputs)[0].flatten()[-32:]

            pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune", fullgraph=True, dynamic=False)
            inputs = self.get_dummy_inputs(torch_device)
            compile_output = pipe(**inputs)[0].flatten()[-32:]

            # Note: Seems to require higher tolerance
            self.assertTrue(np.allclose(normal_output, compile_output, atol=1e-2, rtol=1e-3))
Test that verifies if torch.compile works with torchao quantization.
tests/quantization/torchao/test_torchao.py