Merge pull request #2 from YdaodiG/loss_corrected

fix bug in loss function

Author: judiebig

src/DOSE/learner.py

@@ -138,7 +138,7 @@ def train_step(self, features):
 
     audio = features['clean_speech']
     noisy = features['noisy_speech']
-
+    audio_orig = features['clean_speech'].clone()
 
     N,T= audio.shape
     device = audio.device
@@ -159,7 +159,7 @@ def train_step(self, features):
 
       noisy_audio = noise_scale_sqrt * audio + (1.0 - noise_scale)**0.5 * noise
       predicted = self.model(noisy_audio, t, noisy)
-      loss = self.loss_fn(audio, predicted.squeeze(1))
+      loss = self.loss_fn(audio_orig, predicted.squeeze(1))
 
 
     self.scaler.scale(loss).backward()