(self)
| 130 | return ema_params |
| 131 | |
| 132 | def _load_optimizer_state(self): |
| 133 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint |
| 134 | opt_checkpoint = bf.join( |
| 135 | bf.dirname(main_checkpoint), f"opt{self.resume_step:06}.pt" |
| 136 | ) |
| 137 | if bf.exists(opt_checkpoint): |
| 138 | logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}") |
| 139 | state_dict = dist_util.load_state_dict( |
| 140 | opt_checkpoint, map_location=dist_util.dev() |
| 141 | ) |
| 142 | self.opt.load_state_dict(state_dict) |
| 143 | |
| 144 | def run_loop(self): |
| 145 | while ( |
no test coverage detected