(self, batch, batch_idx, optimizer_idx)
| 148 | return x |
| 149 | |
| 150 | def training_step(self, batch, batch_idx, optimizer_idx): |
| 151 | # https://github.com/pytorch/pytorch/issues/37142 |
| 152 | # try not to fool the heuristics |
| 153 | x = self.get_input(batch, self.image_key) |
| 154 | xrec, qloss, ind = self(x, return_pred_indices=True) |
| 155 | |
| 156 | if optimizer_idx == 0: |
| 157 | # autoencode |
| 158 | aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, |
| 159 | last_layer=self.get_last_layer(), split="train", |
| 160 | predicted_indices=ind) |
| 161 | |
| 162 | self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True) |
| 163 | return aeloss |
| 164 | |
| 165 | if optimizer_idx == 1: |
| 166 | # discriminator |
| 167 | discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, |
| 168 | last_layer=self.get_last_layer(), split="train") |
| 169 | self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True) |
| 170 | return discloss |
| 171 | |
| 172 | def validation_step(self, batch, batch_idx): |
| 173 | log_dict = self._validation_step(batch, batch_idx) |
nothing calls this directly
no test coverage detected