From 2a4dcedbe879ed8ef8ae9d84e255875998c29f11 Mon Sep 17 00:00:00 2001 From: Logan Adams Date: Thu, 14 Nov 2024 15:12:40 -0800 Subject: [PATCH] Formatting --- deepspeed/checkpoint/deepspeed_checkpoint.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/deepspeed/checkpoint/deepspeed_checkpoint.py b/deepspeed/checkpoint/deepspeed_checkpoint.py index 796f5be62bc0..9a368b7a0a25 100644 --- a/deepspeed/checkpoint/deepspeed_checkpoint.py +++ b/deepspeed/checkpoint/deepspeed_checkpoint.py @@ -144,7 +144,10 @@ def get_iteration(self): def get_embedding_state(self, tp_index: int) -> Dict: assert tp_index in self.tp_to_embedding_map.keys() - sd_list = [torch.load(fname, map_location=torch.device('cpu'), weights_only=False) for fname in self.tp_to_embedding_map[tp_index]] + sd_list = [ + torch.load(fname, map_location=torch.device('cpu'), weights_only=False) + for fname in self.tp_to_embedding_map[tp_index] + ] sd = self._merge_state_dicts(sd_list) return sd