def METHOD_NAME(grad: torch.Tensor):
    if len(buffer) == 0:
        buffer.append(torch.zeros_like(grad))
    if self._current_step < self.training_steps:
        buffer[0] += (target.detach() * grad.detach()).pow(2)