Thanks for clarifying the question, as I’ve indeed misunderstood it.
The data augmentation (transformation) will be applied lazily, i.e. while each sample if being loaded.
E.g. if you get the sample at index 0 using x, y = train_dataset[0], the transformations will be applied live at this line of code while executing __getitem__.
The same applies for drawing batches from your DataLoader. While creating the batch, each sample will be drawn from the Dataset by calling its __getitem__, such that the transformation will be applied live again.
Does this answer your question?