(self, beta)
| 6 | |
| 7 | class EMA: |
| 8 | def __init__(self, beta): |
| 9 | super().__init__() |
| 10 | self.beta = beta |
| 11 | self.step = 0 |
| 12 | |
| 13 | def update_model_average(self, ma_model, current_model): |
| 14 | for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()): |