bug fix

realAsma · realAsma · commit 77c48fec063c · 2025-09-15T21:00:52.000Z
Signed-off-by: realAsma &lt;akuriparambi@nvidia.com&gt;
diff --git a/modelopt/torch/quantization/plugins/transformers_trainer.py b/modelopt/torch/quantization/plugins/transformers_trainer.py
@@ -180,7 +180,7 @@ def _save_modelopt_state_with_weights(self):
             if "kd_loss" not in state and "export_student" not in state
         ]
         modelopt_full_state = {
-            "modelopt_state_dict": modelopt_state["modelopt_state_dict"],
+            "modelopt_state": modelopt_state,
             "modelopt_state_weights": get_quantizer_state_dict(self.model),
         }
 
@@ -189,7 +189,7 @@ def _save_modelopt_state_with_weights(self):
 
     def _restore_modelopt_state_with_weights(self):
         modelopt_full_state = torch.load(self._modelopt_state_path, weights_only=False)
-        restore_from_modelopt_state(self.model, modelopt_full_state["modelopt_state_dict"])
+        restore_from_modelopt_state(self.model, modelopt_full_state["modelopt_state"])
         set_quantizer_state_dict(self.model, modelopt_full_state["modelopt_state_weights"])
 
     def _quantize_model(self):

Original file line number	Diff line number	Diff line change
`@@ -180,7 +180,7 @@ def _save_modelopt_state_with_weights(self):`
`180`	`180`	`if "kd_loss" not in state and "export_student" not in state`
`181`	`181`	`]`
`182`	`182`	`modelopt_full_state = {`
`183`		`- "modelopt_state_dict": modelopt_state["modelopt_state_dict"],`
	`183`	`+ "modelopt_state": modelopt_state,`
`184`	`184`	`"modelopt_state_weights": get_quantizer_state_dict(self.model),`
`185`	`185`	`}`
`186`	`186`
`@@ -189,7 +189,7 @@ def _save_modelopt_state_with_weights(self):`
`189`	`189`
`190`	`190`	`def _restore_modelopt_state_with_weights(self):`
`191`	`191`	`modelopt_full_state = torch.load(self._modelopt_state_path, weights_only=False)`
`192`		`- restore_from_modelopt_state(self.model, modelopt_full_state["modelopt_state_dict"])`
	`192`	`+ restore_from_modelopt_state(self.model, modelopt_full_state["modelopt_state"])`
`193`	`193`	`set_quantizer_state_dict(self.model, modelopt_full_state["modelopt_state_weights"])`
`194`	`194`
`195`	`195`	`def _quantize_model(self):`