def on_step_begin(
        self,
        args: TrainingArguments,
        state: TrainerState,
        control: TrainerControl,
        **kwargs,
    ):
        """
        Event called at the beginning of a training step. If using gradient accumulation, one training step might take
        several inputs.
        """
        # memory
        getattr(torch, self.device).reset_peak_memory_stats()

        # time
        self.state.step_start_time = time.perf_counter()

        Event called at the beginning of a training step. If using gradient accumulation, one training step might take
        several inputs.
        
examples/medusa/callback.py