you can use something simple like this:
from torch.utils.data import Dataset as TorchDataset

class SpectrogramDataset(TorchDataset):

    def __init__(self,file_label_ds,  process_func, audio_path=""):
        self.ds= file_label_ds
        self.process_func = process_func
        self.audio_path=audio_path
    def __getitem__(self, index):
        file,label=self.ds[index]
        x=self.process_func(self.audio_path+file)
        return x, file, label
    def __len__(self):
        return len(self.ds)

file_label_ds is a dataset that gives you the file name and label.
process_func is a function that takes the full audio path and return the spectrogram.
PS: this helps you to extract the spectrograms in parallel using the CPU (if you have num_workers>0),  this won’t work if you use the gpu to extract the spectrograms.