diff --git a/codes/data/audio/wav_aug.py b/codes/data/audio/wav_aug.py
index 58e0f2f7..9c8a870d 100644
--- a/codes/data/audio/wav_aug.py
+++ b/codes/data/audio/wav_aug.py
@@ -45,9 +45,9 @@ class WavAugmentor:
         ]
         vol_effect = random.choice(volume_effects)
         effects = [speed_effect, band_effect, vol_effect]
-        out, sr = torchaudio.sox_effects.apply_effects_tensor(wav, sample_rate, effects)
         # Add a variable amount of noise
-        out = out + torch.rand_like(out) * random.random() * .05
+        out = wav + torch.rand_like(wav) * random.random() * .05
+        out, sr = torchaudio.sox_effects.apply_effects_tensor(out, sample_rate, effects)
         return out
 
 
diff --git a/codes/data/audio/wavfile_dataset.py b/codes/data/audio/wavfile_dataset.py
index 9917f83f..06c2913e 100644
--- a/codes/data/audio/wavfile_dataset.py
+++ b/codes/data/audio/wavfile_dataset.py
@@ -63,8 +63,8 @@ class WavfileDataset(torch.utils.data.Dataset):
                 clip2 = self.augmentor.augment(clip2, self.sampling_rate)
 
         return {
-            'clip1': clip1,
-            'clip2': clip2,
+            'clip1': clip1[0, :].unsqueeze(0),
+            'clip2': clip2[0, :].unsqueeze(0),
             'path': filename,
         }
 
@@ -83,10 +83,11 @@ if __name__ == '__main__':
     }
     from data import create_dataset, create_dataloader, util
 
-    ds, c = create_dataset(params, return_collate=True)
+    ds = create_dataset(params, return_collate=True)
     dl = create_dataloader(ds, params, collate_fn=c)
     i = 0
     for b in tqdm(dl):
-        torchaudio.save(f'{i}_clip1.wav', b['clip1'], ds.sampling_rate)
-        torchaudio.save(f'{i}_clip2.wav', b['clip2'], ds.sampling_rate)
-        i += 1
+        for b_ in range(16):
+            torchaudio.save(f'{i}_clip1_{b_}.wav', b['clip1'][b_], ds.sampling_rate)
+            torchaudio.save(f'{i}_clip2_{b_}.wav', b['clip2'][b_], ds.sampling_rate)
+            i += 1