My loader is:

class FeatureLoader(Dataset):
    def __init__(self, data_path, device='cpu'):
        torch.manual_seed(0)

        self.data_path = data_path
        self.device = device
        self.total_length = 0
        self.files = sorted(glob.glob(os.path.join(data_path, "*.pt")))

        self.samples = np.array([])

        for data_file in self.files:
            (audio_pt, feature_pt) = torch.load(data_file, map_location=self.device)

            self.total_length = self.total_length + audio_pt.shape[0]
            self.samples = np.append(self.samples, {"file": data_file, "length": audio_pt.shape[0]})

    def __len__(self):
        return self.total_length

    def __getitem__(self, index):
        samples_index = index % len(self.files)
        sample = self.samples[samples_index]
        (audio_pt, feature_pt) = torch.load(sample['file'], map_location=self.device)

        input_feature = feature_pt[index % sample['length']].detach()
        output_audio = audio_pt[index % sample['length']].detach()

        return input_feature, output_audio

I have 1.5M samples, but my GPU usage doesn’t quite increase at all. Stays low, which means I can probably load more per batch. Currently I have a batch_size of 8192, but if I increase to 16384, then I get some CUDA errors. So I’m wondering what I can do to speed it up.