Fix format

chintu619 · Apr 27, 2022 · 9b33b79 · 9b33b79
1 parent 346a424
commit 9b33b79
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 6 deletions.
diff --git a/espnet2/enh/espnet_model.py b/espnet2/enh/espnet_model.py
@@ -15,6 +15,7 @@
 from espnet2.enh.loss.criterions.time_domain import TimeDomainLoss
 from espnet2.enh.loss.wrappers.abs_wrapper import AbsLossWrapper
 from espnet2.enh.separator.abs_separator import AbsSeparator
+from espnet2.enh.separator.dan_separator import DANSeparator
 from espnet2.torch_utils.device_funcs import force_gatherable
 from espnet2.train.abs_espnet_model import AbsESPnetModel
 
@@ -135,9 +136,11 @@ def forward(
         speech_ref = speech_ref[..., : speech_lengths.max()]
         speech_ref = speech_ref.unbind(dim=1)
         additional = {}
-        additional["feature_ref"] = [
-            self.encoder(r, speech_lengths)[0] for r in speech_ref
-        ]
+        # Additional data is required in Deep Attractor Network
+        if isinstance(self.separator, DANSeparator):
+            additional["feature_ref"] = [
+                self.encoder(r, speech_lengths)[0] for r in speech_ref
+            ]
 
         speech_mix = speech_mix[:, : speech_lengths.max()]
 

diff --git a/espnet2/enh/loss/criterions/tf_domain.py b/espnet2/enh/loss/criterions/tf_domain.py
@@ -346,6 +346,7 @@ def forward(self, ref, inf) -> torch.Tensor:
         Reference:
             Independent Vector Analysis with Deep Neural Network Source Priors;
             Li et al 2020; https://arxiv.org/abs/2008.11273
+
         Args:
             ref: (Batch, T, F) or (Batch, T, C, F)
             inf: (Batch, T, F) or (Batch, T, C, F)

diff --git a/espnet2/tasks/enh.py b/espnet2/tasks/enh.py
@@ -74,18 +74,18 @@
     classes=dict(
         asteroid=AsteroidModel_Converter,
         conformer=ConformerSeparator,
+        dan=DANSeparator,
         dc_crn=DC_CRNSeparator,
         dccrn=DCCRNSeparator,
+        dpcl=DPCLSeparator,
+        dpcl_e2e=DPCLE2ESeparator,
         dprnn=DPRNNSeparator,
         fasnet=FaSNetSeparator,
         rnn=RNNSeparator,
         skim=SkiMSeparator,
         svoice=SVoiceSeparator,
         tcn=TCNSeparator,
         transformer=TransformerSeparator,
-        dpcl=DPCLSeparator,
-        dpcl_e2e=DPCLE2ESeparator,
-        dan=DANSeparator,
         wpe_beamformer=NeuralBeamformer,
     ),
     type_check=AbsSeparator,