device argument for using gpu

melspectrogram.py

@@ -28,13 +28,13 @@ class Standardization(torch.nn.Module):
         return (x - self.mean) / self.std
 
 class MelSpectrogramProcessor:
-    def __init__(self, sample_rate=SR, n_mels=NMELS, n_fft=NFFT, hop_length=HOPLEN, f_min=FMIN, f_max=FMAX):
+    def __init__(self, sample_rate=SR, n_mels=NMELS, n_fft=NFFT, hop_length=HOPLEN, f_min=FMIN, f_max=FMAX, device='cpu'):
         self.transform = nn.Sequential(
             T.MelSpectrogram(sample_rate=sample_rate, n_mels=n_mels, n_fft=n_fft, hop_length=hop_length, f_min=f_min, f_max=f_max),
             T.AmplitudeToDB(),
             Normalization(),
             Standardization(mean=MEAN, std=STD),
-        )
+        ).to(device)
 
     def process(self, waveform):
         return self.transform(waveform)