(self, outputs: EPOCH_OUTPUT, step: str)
| 79 | # ========== ON EPOCH END ========== |
| 80 | |
| 81 | def _shared_epoch_end(self, outputs: EPOCH_OUTPUT, step: str): |
| 82 | assert isinstance(outputs, dict) |
| 83 | with torch.no_grad(): |
| 84 | mean_loss = torch.stack([out[f"{step}/loss"] for out in outputs]).mean() |
| 85 | accuracy = self.__metrics[f"{step}_acc"].compute() |
| 86 | log = {f"{step}/loss": mean_loss, f"{step}/accuracy": accuracy} |
| 87 | self.__metrics[f"{step}_acc"].reset() |
| 88 | self.log_dict(log, on_step=False, on_epoch=True) |
| 89 | |
| 90 | def training_epoch_end(self, outputs: EPOCH_OUTPUT): |
| 91 | self._shared_epoch_end(outputs, "train") |
no outgoing calls