(self, steps_to_save)
| 71 | |
| 72 | class CheckpointingCallback(TrainerCallback): |
| 73 | def __init__(self, steps_to_save): |
| 74 | self.steps_to_save = steps_to_save |
| 75 | |
| 76 | self.rank = torch.distributed.get_rank() |
| 77 | |
| 78 | def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): |
| 79 | if state.global_step in self.steps_to_save: |
nothing calls this directly
no outgoing calls
no test coverage detected