(self, )
| 589 | self.model_ema.reset_num_updates() |
| 590 | |
| 591 | def make_cond_schedule(self, ): |
| 592 | self.cond_ids = torch.full(size=(self.num_timesteps,), fill_value=self.num_timesteps - 1, dtype=torch.long) |
| 593 | ids = torch.round(torch.linspace(0, self.num_timesteps - 1, self.num_timesteps_cond)).long() |
| 594 | self.cond_ids[:self.num_timesteps_cond] = ids |
| 595 | |
| 596 | @rank_zero_only |
| 597 | @torch.no_grad() |