From fbdf76b2fc9c91baefca2d3b1d8a3cd8e2cdd85d Mon Sep 17 00:00:00 2001
From: WeberJulian <julian.weber@hotmail.fr>
Date: Wed, 27 Apr 2022 10:42:26 +0200
Subject: [PATCH] returns y_mask in VITS inference (#1540)

* returns y_mask

* make style
---
 TTS/tts/models/vits.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/TTS/tts/models/vits.py b/TTS/tts/models/vits.py
index 7807efc1..613e4eae 100644
--- a/TTS/tts/models/vits.py
+++ b/TTS/tts/models/vits.py
@@ -1050,7 +1050,15 @@ class Vits(BaseTTS):
 
         o = self.waveform_decoder((z * y_mask)[:, :, : self.max_inference_len], g=g)
 
-        outputs = {"model_outputs": o, "alignments": attn.squeeze(1), "z": z, "z_p": z_p, "m_p": m_p, "logs_p": logs_p}
+        outputs = {
+            "model_outputs": o,
+            "alignments": attn.squeeze(1),
+            "z": z,
+            "z_p": z_p,
+            "m_p": m_p,
+            "logs_p": logs_p,
+            "y_mask": y_mask,
+        }
         return outputs
 
     @torch.no_grad()