(self, batch, batch_idx, optimizer_idx)
| 261 | return x |
| 262 | |
| 263 | def training_step(self, batch, batch_idx, optimizer_idx): |
| 264 | # https://github.com/pytorch/pytorch/issues/37142 |
| 265 | # try not to fool the heuristics |
| 266 | x = self.get_input(batch, self.image_key) |
| 267 | xrec, qloss, ind = self(x, return_pred_indices=True) |
| 268 | |
| 269 | if optimizer_idx == 0: |
| 270 | # autoencode |
| 271 | aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, |
| 272 | last_layer=self.get_last_layer(), split="train", |
| 273 | predicted_indices=ind) |
| 274 | |
| 275 | self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=False, on_epoch=True) |
| 276 | return aeloss |
| 277 | |
| 278 | if optimizer_idx == 1: |
| 279 | # discriminator |
| 280 | discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, |
| 281 | last_layer=self.get_last_layer(), split="train") |
| 282 | self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=False, on_epoch=True) |
| 283 | return discloss |
| 284 | |
| 285 | def validation_step(self, batch, batch_idx): |
| 286 | log_dict = self._validation_step(batch, batch_idx) |
nothing calls this directly
no test coverage detected