Fix return outputs

rhyw · Feb 25, 2022 · 14c1179 · 14c1179
1 parent 424d04e
commit 14c1179
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/TTS/tts/models/vits.py b/TTS/tts/models/vits.py
@@ -1078,7 +1078,7 @@ def train_step(self, batch: dict, criterion: nn.Module, optimizer_idx: int) -> T
                     scores_disc_real,
                     scores_disc_fake,
                 )
-            return {}, loss_dict
+            return outputs, loss_dict
 
         if optimizer_idx == 1:
             mel = batch["mel"]

diff --git a/tests/tts_tests/test_vits.py b/tests/tts_tests/test_vits.py
@@ -410,7 +410,9 @@ def test_train_step(self):
             for _ in range(5):
                 batch = self._create_batch(config, 2)
                 for idx in [0, 1]:
-                    _, loss_dict = model.train_step(batch, criterions, idx)
+                    outputs, loss_dict = model.train_step(batch, criterions, idx)
+                    self.assertFalse(not outputs)
+                    self.assertFalse(not loss_dict)
                     loss_dict["loss"].backward()
                     optimizers[idx].step()
                     optimizers[idx].zero_grad()