Thanks for the reply Patrick. Each numpy array only contains one sample. It turns out DatasetFolder is exactly what I need because I basically needed a generic implementation of ImageFolder. I implemented it like this:
def npy_loader(path):
                sample = torch.from_numpy(np.load(path))
                return sample
train_set = DatasetFolderWithPaths(root=args.data_root, loader=npy_loader, extensions='.npy', transform= transform)