(self)
| 214 | self.mp_trainer.backward(loss) |
| 215 | |
| 216 | def _update_ema(self): |
| 217 | for rate, params in zip(self.ema_rate, self.ema_params): |
| 218 | update_ema(params, self.mp_trainer.master_params, rate=rate) |
| 219 | |
| 220 | def _anneal_lr(self): |
| 221 | if not self.lr_anneal_steps: |