| 330 | print("Replace compute_loss!!") |
| 331 | |
| 332 | class SaveProcessorCallback(TrainerCallback): |
| 333 | def __init__(self, processor): |
| 334 | self.processor = processor |
| 335 | |
| 336 | def on_save(self, args, state, control, **kwargs): |
| 337 | if state.is_world_process_zero: |
| 338 | output_dir = args.output_dir |
| 339 | if state.global_step > 0: |
| 340 | output_dir = os.path.join(args.output_dir, f"checkpoint-{state.global_step}") |
| 341 | self.processor.save_pretrained(output_dir) |
| 342 | return control |
| 343 | |
| 344 | class ProfilerTrainer(Trainer): |
| 345 | def __init__(self, *args, **kwargs): |