Hmm…maybe you are right that even if I use a single neural network inside the Dataset would cause it slower since it is not process in parallel CPUs.
I was thinking to use a single neural in Dataset  and send the it to GPU to augment data, if the neural network is accessible on all worker.  Since GPU would process faster than CPU.
One way to do this is use the class object of collate_fn for the Dataloader, so that the collate_fn object initialize a neural network in GPU to augment in batch. However my dataset is a aiocr dataset which contains character level and sentence level data and I wish to augment in character level.
Maybe we should find a way to do this. Hmm…