From 2989d39239d2f46e584c1e1180ba46b9768afb2a Mon Sep 17 00:00:00 2001 From: hoshi-hiyouga Date: Wed, 23 Apr 2025 16:38:27 +0800 Subject: [PATCH] Merge commit from fork --- scripts/convert_ckpt/llamafy_baichuan2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/convert_ckpt/llamafy_baichuan2.py b/scripts/convert_ckpt/llamafy_baichuan2.py index 3dbeff49..62dc6a51 100644 --- a/scripts/convert_ckpt/llamafy_baichuan2.py +++ b/scripts/convert_ckpt/llamafy_baichuan2.py @@ -32,7 +32,7 @@ def save_weight(input_dir: str, output_dir: str, shard_size: str, save_safetenso baichuan2_state_dict: dict[str, torch.Tensor] = OrderedDict() for filepath in tqdm(os.listdir(input_dir), desc="Load weights"): if os.path.isfile(os.path.join(input_dir, filepath)) and filepath.endswith(".bin"): - shard_weight = torch.load(os.path.join(input_dir, filepath), map_location="cpu") + shard_weight = torch.load(os.path.join(input_dir, filepath), map_location="cpu", weights_only=True) baichuan2_state_dict.update(shard_weight) llama_state_dict: dict[str, torch.Tensor] = OrderedDict()