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

fix test case

Browse Source 
Former-commit-id: c244af0dc3478532de02271667e7af4ad8f54228
This commit is contained in:
hiyouga 2024-06-25 02:51:49 +08:00
parent 0ae1302e41
commit 80effa2993
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tests/model/test_base.py
View File

@ -73,7 +73,8 @@ def test_valuehead():
tokenizer_module["tokenizer"], model_args, finetuning_args, is_trainable=False, add_valuehead=True
)
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained(
ref_model: "AutoModelForCausalLMWithValueHead" = AutoModelForCausalLMWithValueHead.from_pretrained(
TINY_LLAMA_VALUEHEAD, torch_dtype=torch.float16, device_map=get_current_device()
)
ref_model.v_head = ref_model.v_head.to(torch.float16)
compare_model(model, ref_model)