WaveGlow input type fixes (#4151)

Signed-off-by: Jocelyn Huang <[email protected]>
NVIDIA · May 18, 2022 · 76e98dd · 76e98dd
1 parent dc750d6
commit 76e98dd
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/nemo/collections/tts/models/waveglow.py b/nemo/collections/tts/models/waveglow.py
@@ -88,9 +88,7 @@ def convert_spectrogram_to_audio(
     ) -> torch.Tensor:
         with self.nemo_infer():
             self.waveglow.remove_weightnorm()
-            audio = self.waveglow(
-                spec=spec.to(self.waveglow.upsample.weight.dtype), run_inverse=True, audio=None, sigma=sigma
-            )
+            audio = self.waveglow(spec=spec.to(self.waveglow.upsample.weight.dtype), sigma=sigma)
             if denoise:
                 audio = self.denoise(audio=audio, strength=denoiser_strength)
 

diff --git a/nemo/collections/tts/modules/waveglow.py b/nemo/collections/tts/modules/waveglow.py
@@ -134,12 +134,13 @@ def input_types(self):
         if self.mode == OperationMode.infer:
             return {
                 "spec": NeuralType(('B', 'D', 'T'), MelSpectrogramType()),
-                "z": NeuralType(('B', 'D', 'T'), MelSpectrogramType()),
+                "z": NeuralType(('B', 'D', 'T'), MelSpectrogramType(), optional=True),
+                "sigma": NeuralType(optional=True),
             }
         else:
             return {
                 "spec": NeuralType(('B', 'D', 'T'), MelSpectrogramType()),
-                "z": NeuralType(('B', 'D', 'T'), MelSpectrogramType()),
+                "z": NeuralType(('B', 'D', 'T'), MelSpectrogramType(), optional=True),
                 "audio": NeuralType(('B', 'T'), AudioSignal(), optional=True),
                 "run_inverse": NeuralType(elements_type=IntType(), optional=True),
                 "sigma": NeuralType(optional=True),