| 314 | |
| 315 | |
| 316 | class MomentumBuffer: |
| 317 | def __init__(self, momentum: float): |
| 318 | self.momentum = momentum |
| 319 | self.running_average = 0 |
| 320 | |
| 321 | def update(self, update_value: torch.Tensor): |
| 322 | new_average = self.momentum * self.running_average |
| 323 | self.running_average = update_value + new_average |
| 324 | |
| 325 | |
| 326 |