Update dataset.py

dataset.py

@@ -209,19 +209,16 @@ class TrainDataset(Dataset):
 
         sig = sig.reshape(-1).astype(np.float32)
 
-        sig = sig.reshape((1, -1))
         target = torch.tensor(sig.copy())
         p_size = random.choice(self.p_sizes)
 
         sig = np.reshape(sig, (-1, p_size))
         mask = self.mask_generator.gen_mask(len(sig), seed=index)[:, np.newaxis]
         sig *= mask
-        sig = torch.tensor(sig.copy())
+        sig = torch.tensor(sig.copy()).reshape(-1)
 
-        sig = sig.reshape(1, -1)
-
-        target = torch.stft(target.squeeze(0), self.chunk_len, self.stride, window=self.hann,
+        target = torch.stft(target, self.chunk_len, self.stride, window=self.hann,
                             return_complex=False).permute(2, 0, 1).float()
-        sig = torch.stft(sig.squeeze(0), self.chunk_len, self.stride, window=self.hann, return_complex=False)
+        sig = torch.stft(sig, self.chunk_len, self.stride, window=self.hann, return_complex=False)
         sig = sig.permute(2, 0, 1).float()
         return sig, target