PyTorch Lightning Metrics Example

pl_metrics_ex.py

    def training_step(self, batch: DataPoint, batch_idx: int) -> pl.TrainResult:
        planes, target_move, target_val = batch
        pred_move, pred_val = self(planes)
        mse_loss, cross_entropy_loss, loss = self.loss(
            pred_move, pred_val, target_move, target_val
        )
        result = pl.TrainResult(minimize=loss)
        result.log("train_loss", loss, prog_bar=True)
        result.log_dict(
            {
                "train_mse_loss": mse_loss,
                "train_ce_loss": cross_entropy_loss,
                "train_acc": accuracy(pred_move, target_move),
            }
        )
        return result