(self, batch, batch_idx, optimizer_idx)
| 115 | return x |
| 116 | |
| 117 | def training_step(self, batch, batch_idx, optimizer_idx): |
| 118 | inputs = self.get_input(batch, self.image_key) |
| 119 | reconstructions, posterior = self(inputs) |
| 120 | |
| 121 | if optimizer_idx == 0: |
| 122 | # train encoder+decoder+logvar |
| 123 | aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step, |
| 124 | last_layer=self.get_last_layer(), split="train") |
| 125 | self.log("aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True) |
| 126 | self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=False) |
| 127 | return aeloss |
| 128 | |
| 129 | if optimizer_idx == 1: |
| 130 | # train the discriminator |
| 131 | discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step, |
| 132 | last_layer=self.get_last_layer(), split="train") |
| 133 | |
| 134 | self.log("discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True) |
| 135 | self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=False) |
| 136 | return discloss |
| 137 | |
| 138 | def validation_step(self, batch, batch_idx): |
| 139 | log_dict = self._validation_step(batch, batch_idx) |
nothing calls this directly
no test coverage detected