minor

Author: Your Name

modelopt/torch/quantization/utils.py

@@ -461,8 +461,6 @@ def set_quantizer_state_dict(model: nn.Module, quantizer_state_dict: dict):
     from .nn import TensorQuantizer
 
     for name, module in model.named_modules():
-        if (
-            isinstance(module, TensorQuantizer)
-            and get_unwrapped_name(name, model) in quantizer_state_dict
-        ):
-            module.load_state_dict(quantizer_state_dict[get_unwrapped_name(name, model)])
+        key = get_unwrapped_name(name, model)
+        if isinstance(module, TensorQuantizer) and key in quantizer_state_dict:
+            module.load_state_dict(quantizer_state_dict[key])