Update QuantizerBase.load_state_dict function signature

Signed-off-by: Kyunggeun Lee <[email protected]>
Co-authored-by: Kyunggeun Lee <[email protected]>

Author: 2 people

TrainingExtensions/torch/src/python/aimet_torch/v2/quantization/base/quantizer.py

@@ -169,7 +169,7 @@ def state_dict(self, *args, **kwargs):  # pylint: disable=arguments-differ
 
         return state_dict
 
-    def load_state_dict(self, state_dict, strict: bool = True):  # pylint:disable=arguments-differ
+    def load_state_dict(self, state_dict, *args, **kwargs):  # pylint:disable=arguments-differ
         if "_extra_state" not in state_dict:
             is_initialized = OrderedDict(
                 {
@@ -180,7 +180,7 @@ def load_state_dict(self, state_dict, strict: bool = True):  # pylint:disable=ar
             )
             state_dict["_extra_state"] = is_initialized
 
-        ret = super().load_state_dict(state_dict, strict)
+        ret = super().load_state_dict(state_dict, *args, **kwargs)
 
         if version.parse(torch.__version__) < version.parse("1.10"):
             # This is for backward compatibility with torch < 1.10

TrainingExtensions/torch/src/python/aimet_torch/v2/quantization/float/quantizer.py

@@ -194,7 +194,7 @@ def set_extra_state(self, state):
         self.mantissa_bits = state["mantissa_bits"].item()
         super().set_extra_state(state)
 
-    def load_state_dict(self, state_dict, strict: bool = True):
+    def load_state_dict(self, state_dict, *args, **kwargs):
         if "maxval" in state_dict:
             if self.maxval is None:
                 del self.maxval
@@ -203,7 +203,7 @@ def load_state_dict(self, state_dict, strict: bool = True):
             del self.maxval
             self.register_buffer("maxval", None)
 
-        ret = super().load_state_dict(state_dict, strict)
+        ret = super().load_state_dict(state_dict, *args, **kwargs)
         return ret
 
     @property

TrainingExtensions/torch/test/python/v2/quantization/affine/test_affine_quantizer.py

@@ -842,7 +842,8 @@ def test_invalid_encoding_analyzer():
 
 @torch.no_grad()
 @pytest.mark.cuda
-def test_is_initialized(x):
+@pytest.mark.parametrize("assign", [True, False])
+def test_is_initialized(x, assign: bool):
     """
     When: Instantiate a quantizer object
     Then:
@@ -920,7 +921,7 @@ def test_is_initialized(x):
         symmetric=True,
         encoding_analyzer=MinMaxEncodingAnalyzer((10,)),
     )
-    qdq.load_state_dict({"min": -torch.ones(10), "max": torch.ones(10)})
+    qdq.load_state_dict({"min": -torch.ones(10), "max": torch.ones(10)}, assign=assign)
     assert qdq.is_initialized()
 
     """
@@ -933,9 +934,9 @@ def test_is_initialized(x):
         symmetric=True,
         encoding_analyzer=MinMaxEncodingAnalyzer((10,)),
     )
-    qdq.load_state_dict({"min": -torch.ones(10)}, strict=False)
+    qdq.load_state_dict({"min": -torch.ones(10)}, strict=False, assign=assign)
     assert not qdq.is_initialized()  # False; max is not initialized yet
-    qdq.load_state_dict({"max": torch.ones(10)}, strict=False)
+    qdq.load_state_dict({"max": torch.ones(10)}, strict=False, assign=assign)
     assert qdq.is_initialized()
 
     """
@@ -949,17 +950,17 @@ def test_is_initialized(x):
         encoding_analyzer=MinMaxEncodingAnalyzer((10,)),
     )
     uninitialized_state_dict = qdq.state_dict()
-    qdq.load_state_dict(uninitialized_state_dict)
+    qdq.load_state_dict(uninitialized_state_dict, assign=assign)
     assert not qdq.is_initialized()
 
     qdq.min.mul_(1.0)
     partially_initialized_state_dict = qdq.state_dict()
-    qdq.load_state_dict(partially_initialized_state_dict)
+    qdq.load_state_dict(partially_initialized_state_dict, assign=assign)
     assert not qdq.is_initialized()
 
     qdq.max.mul_(1.0)
     fully_initialized_state_dict = qdq.state_dict()
-    qdq.load_state_dict(fully_initialized_state_dict)
+    qdq.load_state_dict(fully_initialized_state_dict, assign=assign)
     assert qdq.is_initialized()
 
     """
@@ -981,7 +982,7 @@ def test_is_initialized(x):
         symmetric=True,
         encoding_analyzer=MinMaxEncodingAnalyzer((10,)),
     )
-    qdq.load_state_dict({"min": -torch.ones(10), "max": torch.ones(10)})
+    qdq.load_state_dict({"min": -torch.ones(10), "max": torch.ones(10)}, assign=assign)
     qdq = copy.deepcopy(qdq)
     assert qdq.is_initialized()
 
@@ -1005,7 +1006,7 @@ def test_is_initialized(x):
         symmetric=True,
         encoding_analyzer=MinMaxEncodingAnalyzer((10,)),
     )
-    qdq.load_state_dict({"min": -torch.ones(10), "max": torch.ones(10)})
+    qdq.load_state_dict({"min": -torch.ones(10), "max": torch.ones(10)}, assign=assign)
     out_before = qdq(x.view(-1, 10))
     res = pickle.dumps(qdq)
     qdq = pickle.loads(res)