allow for a custom reconstruction target in soundstream

lucidrains · Apr 28, 2023 · d05c020 · d05c020
1 parent c3fabc3
commit d05c020
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/audiolm_pytorch/soundstream.py b/audiolm_pytorch/soundstream.py
@@ -651,6 +651,7 @@ def seq_len_multiple_of(self):
     def forward(
         self,
         x,
+        target = None,
         return_encoded = False,
         return_discr_loss = False,
         return_discr_losses_separately = False,
@@ -752,7 +753,9 @@ def forward(
 
         # recon loss
 
-        recon_loss = F.mse_loss(orig_x, recon_x)
+        target = default(target, orig_x)  # target can also be passed in, in the case of denoising
+
+        recon_loss = F.mse_loss(target, recon_x)
 
         # multispectral recon loss - eq (4) and (5) in https://arxiv.org/abs/2107.03312
 

diff --git a/audiolm_pytorch/version.py b/audiolm_pytorch/version.py
@@ -1 +1 @@
-__version__ = '0.28.1'
+__version__ = '0.28.2'