luvwinnie:

I have done this because that my number of data is too large and it take too much time to use the neural network(network such as StyleGAN) to augment the dataset.


I’m unsure how to understand this in the context of the question.
If the single neural network is too slow in its performance to transform the data fast enough, I would expect that model clones would be desired in each worker, so that the processing could be done in parallel (on the CPU). However, it seems you would like to use a single neural network inside the Dataset (and all workers), which would then yield the slow performance again, wouldn’t it?