(self, module)
| 16 | self.shadow[name] = param.data.clone() |
| 17 | |
| 18 | def update(self, module): |
| 19 | if isinstance(module, nn.DataParallel): |
| 20 | module = module.module |
| 21 | for name, param in module.named_parameters(): |
| 22 | if param.requires_grad: |
| 23 | self.shadow[name].data = ( |
| 24 | 1. - |
| 25 | self.mu) * param.data + self.mu * self.shadow[name].data |
| 26 | |
| 27 | def ema(self, module): |
| 28 | if isinstance(module, nn.DataParallel): |