Hi, @rwightman, thanks a lot for your answer and kind suggestions!
Yeah, the train split has ten numpy file, the validation split has two numpy file, and the test split has two numpy file.
Each numpy file has a numpy dictionary contains 10000 samples, with two keys: batch_x  with size 10000x512x64 and batch_y with size 10000x256.
What you suggested above is too high-level for me, especially the caching-related.
In fact, my concern is on the batch inputs. If we set the batch size as 256, and the sample indices are randomly selected from 0~10*10000, then when the Dataset wants to generate this batch, he may need to read all ten numpy files and then get the corresponding samples to fill the batch. When using cache, the data in all ten numpy files will be cached, I think there is no difference between this kind of caching and reading all data into memory just once (like the way in https://github.com/pytorch/vision/blob/master/torchvision/datasets/cifar.py 38). Maybe I’m misunderstanding on caching.