Failed to save the static quantized model

[yiliu30]

Hi, I am following this example and want to save the INT8 static quantization result, but it’s failing.
Could you take a look, thanks!

    ... 
    # quantized linear represented as an nn.Linear with modified tensor subclass weights
    # for both activation and weight quantization
    quantize_(m, ApplyStaticQuantConfig(target_dtype), is_observed_linear)
    print("quantized model (applying tensor subclass to weight):", m)
    after_quant = m(*example_inputs)  # <---- Line 310
    torch.save(m.state_dict(), f"qmodel_m.pt")  # <---- My change
    assert compute_error(before_quant, after_quant) > 25
    print("test passed")

• Log

Testing torch.uint8 static quantization:
example inputs shape: torch.Size([1, 64])
quantized model (applying tensor subclass to weight): ToyLinearModel(
  (linear1): Linear(in_features=64, out_features=64, bias=False)
  (linear2): Linear(in_features=64, out_features=32, bias=False)
)
Traceback (most recent call last):
  File "/home/user/workspace/torchao/tutorials/calibration_flow/static_quant.py", line 325, in <module>
    test_static_quant(torch.uint8, MappingType.ASYMMETRIC)
  File "/home/user/workspace/torchao/tutorials/calibration_flow/static_quant.py", line 311, in test_static_quant
    torch.save(m.state_dict(), f"qmodel_m.pt")
  File "/home/user/miniforge3/envs/ao/lib/python3.11/site-packages/torch/serialization.py", line 965, in save
    _save(
  File "/home/user/miniforge3/envs/ao/lib/python3.11/site-packages/torch/serialization.py", line 1211, in _save
    pickler.dump(obj)
AttributeError: Can't pickle local object '_apply_static_quant_transform.<locals>.<lambda>'

• Env info
• torchao 0.10.0+gita99598d8
• torch 2.8.0.dev20250324+cu126

cc @jerryzh168