(self, batch, batch_idx)
| 597 | @torch.no_grad() |
| 598 | # def on_train_batch_start(self, batch, batch_idx, dataloader_idx): |
| 599 | def on_train_batch_start(self, batch, batch_idx): |
| 600 | # only for very first batch |
| 601 | if self.scale_by_std and self.current_epoch == 0 and self.global_step == 0 and batch_idx == 0 and not self.restarted_from_ckpt: |
| 602 | assert self.scale_factor == 1., 'rather not use custom rescaling and std-rescaling simultaneously' |
| 603 | # set rescale weight to 1./std of encodings |
| 604 | print("### USING STD-RESCALING ###") |
| 605 | x = super().get_input(batch, self.first_stage_key) |
| 606 | x = x.to(self.device) |
| 607 | encoder_posterior = self.encode_first_stage(x) |
| 608 | z = self.get_first_stage_encoding(encoder_posterior).detach() |
| 609 | del self.scale_factor |
| 610 | self.register_buffer('scale_factor', 1. / z.flatten().std()) |
| 611 | print(f"setting self.scale_factor to {self.scale_factor}") |
| 612 | print("### USING STD-RESCALING ###") |
| 613 | |
| 614 | def register_schedule(self, |
| 615 | given_betas=None, beta_schedule="linear", timesteps=1000, |
nothing calls this directly
no test coverage detected