| 373 | return ema_model |
| 374 | |
| 375 | def save_pretrained(self, path): |
| 376 | if self.model_cls is None: |
| 377 | raise ValueError("`save_pretrained` can only be used if `model_cls` was defined at __init__.") |
| 378 | |
| 379 | if self.model_config is None: |
| 380 | raise ValueError("`save_pretrained` can only be used if `model_config` was defined at __init__.") |
| 381 | |
| 382 | model = self.model_cls.from_config(self.model_config) |
| 383 | state_dict = self.state_dict() |
| 384 | state_dict.pop("shadow_params", None) |
| 385 | |
| 386 | model.register_to_config(**state_dict) |
| 387 | self.copy_to(model.parameters()) |
| 388 | model.save_pretrained(path) |
| 389 | |
| 390 | def get_decay(self, optimization_step: int) -> float: |
| 391 | """ |