How to collect batched predictions?

[VietTralala]

Hello :)
Currently I use trainer.predict(model=..., dataloaders=...) which returns the results of predict_step(...) in a list where each element in the list corresponds to one batch input to the predict_step function which I already implemented. I am looking for an predict_epoch_end kind of function to collect to batched predictions into one data structure but only found the possibility to define a callback on_prediction_epoch_end() but this has no return possibility. How should one best proceed to collect the batched predictions?

i.e. each predict_step() returns a tensor of shape batch_size x 10
and lets assume the dataloader gives 5 batches to predict. Then the trainer.predict() function will return a list of lenght 5 with each element beeing a tensor of shape batch_size x 10. However I would rather like to receive one tensor of shape 5*batch_size x 10.

This is just a simple example to illustrate my problem. My actual return per prediction step is a little bit more involved and I would like to clean up the structure before returning it and also make this cleanup logic part of the module so that I dont have to remember the specifics.

Is there a way to do that in the current framework and if so how?
Thanks in advance for any help!

[rohitgr7]

Issue to track #9380