From c2b50a6daac243c5cc5265ea9548153344b3fe83 Mon Sep 17 00:00:00 2001 From: Will Johnson Date: Tue, 10 Dec 2024 11:32:22 -0500 Subject: [PATCH] fmt Signed-off-by: Will Johnson --- .../src/fms_acceleration_moe/utils/checkpoint_utils.py | 4 +++- 1 file changed, 3 insertions(+), 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 index c43ba1c2..fb8ab1bc 100644 --- a/plugins/accelerated-moe/src/fms_acceleration_moe/utils/checkpoint_utils.py +++ b/plugins/accelerated-moe/src/fms_acceleration_moe/utils/checkpoint_utils.py @@ -433,7 +433,9 @@ def save_sharded_safetensors( ".safetensors", "{suffix}.safetensors" ) state_dict_split = split_torch_state_dict_into_shards( - input_state_dict, filename_pattern=filename_pattern, max_shard_size=max_shard_size + input_state_dict, + filename_pattern=filename_pattern, + max_shard_size=max_shard_size, ) index = { "metadata": state_dict_split.metadata,