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

fix #6546

Browse Source
This commit is contained in:
hiyouga
2025-01-07 06:30:44 +00:00
parent b832ed9a60
commit 870f23d7ea
3 changed files with 32 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/llamafactory/train/kto/trainer.py
View File

@@ -30,7 +30,7 @@ from typing_extensions import override
from ...extras.constants import IGNORE_INDEX
from ...extras.packages import is_transformers_version_equal_to_4_46, is_transformers_version_greater_than
from ..callbacks import SaveProcessorCallback
from ..trainer_utils import create_custom_optimizer, create_custom_scheduler, get_batch_logps
from ..trainer_utils import create_custom_optimizer, create_custom_scheduler, get_batch_logps, nested_detach
if TYPE_CHECKING:
@@ -142,7 +142,7 @@ class CustomKTOTrainer(KTOTrainer):
r"""
Runs forward pass and computes the log probabilities.
"""
batch = {k: v.detach().clone() for k, v in batch.items()} # avoid error
batch = nested_detach(batch, clone=True) # avoid error
model_inputs = {
"input_ids": batch[f"{prefix}input_ids"],
"attention_mask": batch[f"{prefix}attention_mask"],