You should consider using torch.utils.data.DataLoader, and specify number of workers. These workers retrieve data from the dataset and will significantly improve the read spead. Here is a little snippet:
dataloader = torch.utils.data.DataLoader(dataset,
 batch_size=512,
 shuffle=True,
num_workers=10)