ros/LLaMA-Factory
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

tiny fix

Browse Source 
Former-commit-id: f7f440986b0ae3b38ea9f2da80789629d4f79ea1
This commit is contained in:
hiyouga
2024-06-16 01:06:41 +08:00
parent 14f7bfc545
commit 05f3a3c944
22 changed files with 27 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,6 +1,6 @@
# Copyright 2024 HuggingFace Inc. and the LlamaFactory team.
#
# This code is inspired by HuggingFace's TRL library.
# This code is inspired by the HuggingFace's TRL library.
# https://github.com/huggingface/trl/blob/v0.8.0/trl/trainer/kto_trainer.py
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -114,8 +114,8 @@ class CustomKTOTrainer(KTOTrainer):
def _save(self, output_dir: Optional[str] = None, state_dict: Optional[Dict[str, "torch.Tensor"]] = None) -> None:
super()._save(output_dir, state_dict)
output_dir = output_dir if output_dir is not None else self.args.output_dir
if self.processor is not None:
output_dir = output_dir if output_dir is not None else self.args.output_dir
getattr(self.processor, "image_processor").save_pretrained(output_dir)
def forward(