def test_memory_footprint(self):
        r"""
        A simple test to check if the model conversion has been done correctly by checking on the
        memory footprint of the converted model and the class type of the linear layers of the converted models
        """
        for model_id in ["hf-internal-testing/tiny-flux-pipe", "hf-internal-testing/tiny-flux-sharded"]:
            transformer_int4wo = self.get_dummy_components(TorchAoConfig("int4wo"), model_id=model_id)["transformer"]
            transformer_int4wo_gs32 = self.get_dummy_components(
                TorchAoConfig("int4wo", group_size=32), model_id=model_id
            )["transformer"]
            transformer_int8wo = self.get_dummy_components(TorchAoConfig("int8wo"), model_id=model_id)["transformer"]
            transformer_bf16 = self.get_dummy_components(None, model_id=model_id)["transformer"]

            # Will not quantized all the layers by default due to the model weights shapes not being divisible by group_size=64
            for block in transformer_int4wo.transformer_blocks:
                self.assertTrue(isinstance(block.ff.net[2].weight, AffineQuantizedTensor))
                self.assertTrue(isinstance(block.ff_context.net[2].weight, AffineQuantizedTensor))

            # Will quantize all the linear layers except x_embedder
            for name, module in transformer_int4wo_gs32.named_modules():
                if isinstance(module, nn.Linear) and name not in ["x_embedder"]:
                    self.assertTrue(isinstance(module.weight, AffineQuantizedTensor))

            # Will quantize all the linear layers
            for module in transformer_int8wo.modules():
                if isinstance(module, nn.Linear):
                    self.assertTrue(isinstance(module.weight, AffineQuantizedTensor))

            total_int4wo = get_model_size_in_bytes(transformer_int4wo)
            total_int4wo_gs32 = get_model_size_in_bytes(transformer_int4wo_gs32)
            total_int8wo = get_model_size_in_bytes(transformer_int8wo)
            total_bf16 = get_model_size_in_bytes(transformer_bf16)

            # TODO: refactor to align with other quantization tests
            # Latter has smaller group size, so more groups -> more scales and zero points
            self.assertTrue(total_int4wo < total_int4wo_gs32)
            # int8 quantizes more layers compare to int4 with default group size
            self.assertTrue(total_int8wo < total_int4wo)
            # int4wo does not quantize too many layers because of default group size, but for the layers it does
            # there is additional overhead of scales and zero points
            self.assertTrue(total_bf16 < total_int4wo)

        A simple test to check if the model conversion has been done correctly by checking on the
        memory footprint of the converted model and the class type of the linear layers of the converted models
        
tests/quantization/torchao/test_torchao.py