Trainer: how to apply a method on a dataloader

[quentinblampey]

I have a method that creates a tensor based on one mini-batch of data, let's call it model.transform.
I want to transform every batch of a dataloader named loader and concatenate all the resulting tensor (with a progress bar showing how many batch we have processed).

Of course, I could easily do it myself, but I'm looking for a one-liner that uses pytorch lightning.
For that, I could use trainer.predict and implement model.predict_step as described on the first line. It does exactly what I want (well, actually, except for the concatenation). The real problem is that I want to be able to use different transformations, but we only have one model.predict_step function.