(self, model, update_fn)
| 468 | self.module.to(device=device) |
| 469 | |
| 470 | def _update(self, model, update_fn): |
| 471 | with torch.no_grad(): |
| 472 | for ema_v, model_v in zip(self.module.state_dict().values(), |
| 473 | model.state_dict().values()): |
| 474 | if self.device is not None: |
| 475 | model_v = model_v.to(device=self.device) |
| 476 | ema_v.copy_(update_fn(ema_v, model_v)) |
| 477 | |
| 478 | def update(self, model): |
| 479 | self._update(model, |