Make validation_step return outputs

[Hravan]

I'd like to create my own callback that will create an image of a confusion matrices for multilabel classification in TensorBoard. To do that at the end of validation epoch, I have to somehow get to outputs of all validation steps. I am trying to implement on_validation_batch end in my custom callback to store outputs internally after every validation batch:

class ConfusionMatrixCallback(pl.Callback):
    def __init__(self, num_classes: int) -> None:
        super().__init__()
        self.num_classes = num_classes
        self.outputs = []

    def on_validation_batch_end(
        self,
        trainer: "pl.Trainer",
        pl_module: "pl.LightningModule",
        outputs,
        batch,
        batch_idx: int,
        dataloader_idx: int,
    ) -> None:
        self.outputs.append(outputs)

The problem is outputs are None, which I guess is because ImageClassifier's validation_step doesn't return anything. How can I easily override it? Of course I can make my own LightningModule, but if I were to, I could as well not use Lightning Flash at all ;) I went through the code and I see this validation_step to be overriden is a few levels deep in an adapter that is an attribute of ImageClassifier instance, which makes it hard to override.

Of course, if you have any suggestion how I can implement what I want differently, I am open to it.