def test_linear_quantizer_preserved_attributes(algorithm, model_dict):
    """
    Test if methods and attributes on the model are preserved by passing
    preserved_attributes to the config.
    """

    class MyModel(nn.Sequential):
        def __init__(self, model_dict):
            super().__init__(model_dict)
            self.conv.weight.data = torch.ones_like(self.conv.weight.data)

        def my_method(self):
            return self.weight + torch.ones_like(self.weight)

        @property
        def weight(self):
            return (
                self.conv.weight
                if hasattr(self.conv, "weight")
                else self.conv.get_submodule("0").weight
            )

    preserved_attrs = ["key_1", "key_2", "my_method", "weight"]

    model = MyModel(model_dict)
    model.key_1 = 5
    model.key_2 = torch.tensor(5)

    config = LinearQuantizerConfig.from_dict(
        {
            "global_config": {
                "algorithm": algorithm,
                "milestones": [0, 3, 4, 5],
            },
            "preserved_attributes": preserved_attrs,
        }
    )
    quantizer_1 = LinearQuantizer(model, LinearQuantizerConfig())
    prepared_model = quantizer_1.prepare(example_inputs=(torch.randn(1, 1, 28, 28),), inplace=False)
    for attr in preserved_attrs:
        assert not hasattr(prepared_model, attr)

    quantizer_2 = LinearQuantizer(model, config)
    prepared_model = quantizer_2.prepare(example_inputs=(torch.randn(1, 1, 28, 28),), inplace=False)
    for attr in preserved_attrs:
        assert hasattr(prepared_model, attr)
    assert torch.all(
        prepared_model.my_method() == 2 * torch.ones_like(prepared_model.conv.weight.data)
    )

    quantizer_2.step()
    prepared_model(torch.randn(2, 1, 28, 28))
    final_model = quantizer_2.finalize()
    for attr in preserved_attrs:
        assert hasattr(final_model, attr)
    assert torch.all(
        final_model.my_method()
        == final_model.weight.data + torch.ones_like(prepared_model.weight.data)
    )

    Test if methods and attributes on the model are preserved by passing
    preserved_attributes to the config.
    
coremltools/test/optimize/torch/quantization/test_quantizer.py