fix: nit return sd (#131)

Signed-off-by: Mehant Kammakomati <mehant.kammakomati2@ibm.com>
foundation-model-stack · Feb 27, 2025 · 706d5cc · 706d5cc
1 parent c959c6b
commit 706d5cc
Showing 1 changed file with 1 addition and 1 deletion.
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