(self)
| 139 | return ema_params |
| 140 | |
| 141 | def _load_optimizer_state(self): |
| 142 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint |
| 143 | opt_checkpoint = bf.join( |
| 144 | bf.dirname(main_checkpoint), f"opt{self.resume_step:06}.pt" |
| 145 | ) |
| 146 | if bf.exists(opt_checkpoint): |
| 147 | logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}") |
| 148 | state_dict = dist_util.load_state_dict( |
| 149 | opt_checkpoint, map_location=dist_util.dev() |
| 150 | ) |
| 151 | self.opt.load_state_dict(state_dict) |
| 152 | |
| 153 | def run_loop(self): |
| 154 | while ( |
no test coverage detected