423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-07-31 10:42:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge commit from fork

Browse Source
This commit is contained in:
Yaowei Zheng 2025-06-26 13:55:42 +08:00 committed by GitHub
parent 7242caf0ff
commit bb7bf51554
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/llamafactory/model/model_utils/valuehead.py
View File

@ -49,7 +49,7 @@ def load_valuehead_params(path_or_repo_id: str, model_args: "ModelArguments") ->
try:
vhead_file = cached_file(filename=V_HEAD_WEIGHTS_NAME, **kwargs)
return torch.load(vhead_file, map_location="cpu")
return torch.load(vhead_file, map_location="cpu", weights_only=True)
except Exception as err:
err_text = str(err)

2
src/llamafactory/train/callbacks.py
View File

@ -76,7 +76,7 @@ def fix_valuehead_checkpoint(
state_dict: dict[str, torch.Tensor] = {key: f.get_tensor(key) for key in f.keys()}
else:
path_to_checkpoint = os.path.join(output_dir, WEIGHTS_NAME)
state_dict: dict[str, torch.Tensor] = torch.load(path_to_checkpoint, map_location="cpu")
state_dict: dict[str, torch.Tensor] = torch.load(path_to_checkpoint, map_location="cpu", weights_only=True)
os.remove(path_to_checkpoint)
decoder_state_dict, v_head_state_dict = {}, {}