fix: nit return sd (#131)

kmehant · web-flow · commit 706d5cc6ffad · 2025-02-27T18:32:11.000+08:00
Signed-off-by: Mehant Kammakomati &lt;mehant.kammakomati2@ibm.com&gt;
diff --git a/plugins/accelerated-moe/src/fms_acceleration_moe/utils/checkpoint_utils.py b/plugins/accelerated-moe/src/fms_acceleration_moe/utils/checkpoint_utils.py
@@ -235,7 +235,7 @@ def get_state_dict_from_dcp_checkpoint(
         planner=_EmptyStateDictLoadPlanner(),
         no_dist=True,
     )
-    return [KEY_MODEL]
+    return sd[KEY_MODEL]
 
 
 # function to get state dict from regular checkoint

Original file line number	Diff line number	Diff line change
`@@ -235,7 +235,7 @@ def get_state_dict_from_dcp_checkpoint(`
`235`	`235`	`planner=_EmptyStateDictLoadPlanner(),`
`236`	`236`	`no_dist=True,`
`237`	`237`	`)`
`238`		`- return [KEY_MODEL]`
	`238`	`+ return sd[KEY_MODEL]`
`239`	`239`
`240`	`240`
`241`	`241`	`# function to get state dict from regular checkoint`