Torch geometric - modify the graphs nodes on the fly

Hi I was wondering if it was possible for my data augmentation to be done on the fly during my epoch, this way it would be convenient to store less data.
I created a function let’s say f that does a slight modification of a graph in input (Data(nodes, edges, edges_attributes, y))
when I charge my batch let’s say hypothetically 5 x1 … x5 I would instead input f(x1) … f(x5) but I noticed that with torch geometric the loading of a batch is kinda a black box and is loading different batches.