Former-commit-id: c36c5c61fc022b3f144d4c798ec584c4954b0181
This commit is contained in:
hiyouga
2024-10-29 13:02:13 +00:00
parent 22912eba1a
commit 43bd1b070c
5 changed files with 16 additions and 14 deletions

View File

@@ -124,12 +124,12 @@ class SaveProcessorCallback(TrainerCallback):
def on_save(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
if args.should_save:
output_dir = os.path.join(args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}")
getattr(self.processor, "image_processor").save_pretrained(output_dir)
self.processor.save_pretrained(output_dir)
@override
def on_train_end(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
if args.should_save:
getattr(self.processor, "image_processor").save_pretrained(args.output_dir)
self.processor.save_pretrained(args.output_dir)
class PissaConvertCallback(TrainerCallback):