Hi ! I think you can use the PyTorch DataLoader:

with num_workers>0, you can use multiprocessing to load your data in parallel to the optimizations steps
with the prefetch_factor  parameter, you can say how many samples are loaded in advance by each data loading worker.

Let me know if that was what you were looking for