code: def test_serialize_quantized_mlp(weights, dtype, device):
      if dtype == torch.float16 and device.type == "cpu":
          pytest.skip(
              "Matrix multiplication is not supported for float16 on CPU"
          )
      input_features = 32
      hidden_features = 10
      output_features = 128
      model = (
          MLP(input_features, hidden_features, output_features)
          .to(dtype)
          .to(device)
      )
      quantize(model, weights=weights)
      inputs = random_tensor((1, 10, input_features), dtype=dtype).to(
          device
      )
      with Calibration():
          model(inputs)
      freeze(model)
      b = io.BytesIO()
      torch.save(model.state_dict(), b)
      b.seek(0)
      state_dict = torch.load(b)
      model_reloaded = MLP(
          input_features, hidden_features, output_features
      ).to(device)
      quantize(model_reloaded)
      # When reloading we must assign instead of copying to force quantized tensors assignment
      model_reloaded.load_state_dict(state_dict, assign=True)
      for name, module in model.named_modules():
          if isinstance(module, QModuleMixin):
              module_reloaded = getattr(model_reloaded, name)
              assert torch.equal(
                  module_reloaded.weight._data, module.weight._data
              )
              assert torch.equal(
                  module_reloaded.weight._scale, module.weight._scale
              )
              assert torch.equal(
                  module_reloaded.input_scale, module.input_scale
              )
              assert torch.equal(
                  module_reloaded.output_scale, module.output_scale
              )
  
file_path: quanto/test/test_serialization.py
imports: import io
  import pytest
  import torch
  from helpers import random_qtensor, random_tensor
  from quanto import Calibration, QTensor, absmax_scale, freeze, quantize
  from quanto.nn import QLinear, QModuleMixin
label_desc: Use image node_id 3 for calling a global function with example usage: test_serialize_quantized_mlp(weights, dtype, device) without return types