Enabling batch transforms
🚀 Feature
Hi ! Well the feature is quite simple, it is to have a simple way to code batch level transforms without having to put them directly in training_step.
Motivation
Motivation comes from the fact that in a project of mine, I want to change the color space of my input images directly on the batch rather than on individual images so that I gain some time. Problem is I didn't find a way to automate that using lightning. When going through code and documentation, I even stumbled upon something that confused me a lot : in the documentation for LightningDataModule, there are numerous examples of places where transforms are passed directly to a pytorch DataLoader constructor. However we can't do that, there is no transform argument in pytorch DataLoader. I looked for an override somewhere in lightning's source code, and tried to find if there is somewhere in the code where the batch could get modified based on transforms but I couldn't find any. So please if I missed something just tell me and I guess this issue can be closed. Else I'm pretty sure this feature could come in handy for others than me.
Pitch
The idea is to allow the user to specify batch transforms in a LightningDataModule or somewhere in the trainer and have them be called automatically after the dataloader is iterated through. The idea for me is to having it be called right after the batch is transferred to GPU in training_forward. I'm pretty sure further details need to be discussed as the implementation will depend on the coding philosophy of the library.
Alternatives
For now the alternative I use is to manually insert he batch transform in training_step. This could also be automated by adding a hook that allows batch modification right before training_step. It would allow for a broader use and could be what is used behind the scene fore a higher level batch transform feature. Again, the details probably need to be discussed.
Additional context
I'll add a screenshot of the confusing part of the documentation I was refering to earlier.

I'll also add that I am willing to work on this feature myself once details are set.