How to get predictions after evaluation phase in a callback

I want to log prediction distribution after each evaluation phase to Weights & Biases. I thought of adding a callback to do this, but couldn’t find anything which allows me to access predictions in on_evaluate method to do so. Is there any way to achieve this or am I missing something ?

If you want to get the predictions back, you should use Trainer.predict().