Extend DataLoader with transform arguments to keep that logic in the DataModule
🚀 Feature

Extend DataLoader to accept transform.
Possibly split into 2 arguments:

sample_transform
batch_transform

The new process of data retrieval would go as follow:

Retrieve sample tuple with (input, target)
sample_transform(sample)
Concatenate to batch
batch_transform(batch)

Motivation
The new DataModule page of the documentation incorrectly includes a transform=transforms argument for the DataLoader:
class MNISTDataModule(pl.LightningDataModule):
    def train_dataloader(self):
        transforms = transform_lib.Compose([
            transform_lib.ToTensor(),
            transform_lib.Normalize(mean=(0.5,), std=(0.5,)),
        ])
        return DataLoader(self.train_dataset, transform=transforms, batch_size=64)
This extra argument would allow to truly separate the Dataset and DataModule in terms of functionality.
If we do transforms on the samples after retrieving it from the Dataset, we could also logically include not only sample transforms, but also possible batch transforms.
Pitch
Dataset and DataModule are not truly separated in functionality yet.
While the documentation suggests that transforms should be specified in DataModule, in the end the transform is still passed on to Dataset.
This requires Dataset to be transform-aware. If a Dataset is not configured with something similar to:
if self.transform:
    input = self.transform(input)
, it will not be possible to transform the data.
If it is necessary to write transform code in the Dataset class, why not specify it directly there?
Extra 1: If Dataset does not have to be tuned to expect PyTorch transforms, this Dataset could be used for both PyTorch and scikit-learn without encoding specific logic for this in the Dataset class.
Extra 2: If we move transform to DataLoader, we can also both specify per-sample transforms and per-batch transforms.
p.s. Would it be faster to transform a batch of numpy arrays into Tensor, or is there no different with transforming the arrays 1-by-1 into a Tensor first before batching them?
Alternatives
Keep it as-is, passing the transform to a Dataset and require the user to make sure the Dataset can handle potential transforms.
Additional context
Liked by: @Borda
Others mentioned by him: @nateraw @awaelchli @williamFalcon