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