Fix

cjluo-nv · cjluo-nv · commit 219bb2d7b134 · 2026-02-09T19:31:53.000Z
Signed-off-by: Chenjie Luo &lt;chenjiel@nvidia.com&gt;
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
@@ -1015,7 +1015,7 @@ def export_hf_checkpoint(
         # Save model
         model.save_pretrained(
             export_dir,
-            state_dict=post_state_dict | extra_state_dict,
+            state_dict={**post_state_dict, **(extra_state_dict or {})},
             save_modelopt_state=save_modelopt_state,
         )
 

Original file line number	Diff line number	Diff line change
`@@ -1015,7 +1015,7 @@ def export_hf_checkpoint(`
`1015`	`1015`	`# Save model`
`1016`	`1016`	`model.save_pretrained(`
`1017`	`1017`	`export_dir,`
`1018`		`- state_dict=post_state_dict \| extra_state_dict,`
	`1018`	`+ state_dict={post_state_dict, (extra_state_dict or {})},`
`1019`	`1019`	`save_modelopt_state=save_modelopt_state,`
`1020`	`1020`	`)`
`1021`	`1021`