Bug in torch.ao.nn.quantized.Sigmoid Parameter Restoration after state_dict Loading #147818
Description
🐛 Describe the bug
There seems to be an issue in PyTorch's quantized Sigmoid module (nnq_Sigmoid) where the quantization parameters (scale and zero_point) are not properly restored when loading the state dictionary (state_dict) into a newly initialized module with different initial parameters.
Code:
import torch
from torch.ao.nn.quantized import Sigmoid as nnq_Sigmoid
def test_sigmoid_serialization():
# Original parameters
scale_original = 0.1
zero_point_original = 5
# Create original module and save state
quant_mod_original = nnq_Sigmoid(scale_original, zero_point_original)
state_dict = quant_mod_original.state_dict()
# New parameters (different from original)
scale_new = 0.5
zero_point_new = 10
# Create new module and load original state
quant_mod_new = nnq_Sigmoid(scale_new, zero_point_new)
quant_mod_new.load_state_dict(state_dict)
# Check if parameters were restored
print("quant_mod_new.output_scale:", quant_mod_new.output_scale)
print("scale_original: ", scale_original)
print("quant_mod_new.output_zero_point:", quant_mod_new.output_zero_point)
print("zero_point_original:", zero_point_original)
test_sigmoid_serialization()
Output:
The parameters scale and zero_point are not restored correctly after loading the state dictionary. The output shows that the parameters are not matching the original values, which implies that the state dictionary is not correctly restoring the quantization parameters.
quant_mod_new.output_scale: 0.5
scale_original: 0.1
quant_mod_new.output_zero_point: 10
zero_point_original: 5
Versions
PyTorch version: 2.6.0+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A
cc @jerryzh168 @jianyuh @raghuramank100 @jamesr66a @vkuzo @jgong5 @Xia-Weiwen @leslie-fang-intel @msaroufim