restored flatten_parameters

Signed-off-by: Boris Fomitchev <bfomitchev@nvidia.com>
NVIDIA · Nov 10, 2022 · d48393c · d48393c
1 parent da91270
commit d48393c
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/nemo/collections/tts/modules/common.py b/nemo/collections/tts/modules/common.py
@@ -123,15 +123,15 @@ def lstm_tensor(self, context: Tensor, lens: Tensor, enforce_sorted: bool = Fals
         seq = nn.utils.rnn.pack_padded_sequence(
             context, lens.long().cpu(), batch_first=True, enforce_sorted=enforce_sorted
         )
-        # if not torch.jit.is_scripting():
-        #    self.bilstm.flatten_parameters()
+        if not torch.jit.is_scripting():
+            self.bilstm.flatten_parameters()
         ret, _ = self.bilstm(seq)
         return nn.utils.rnn.pad_packed_sequence(ret, batch_first=True)
 
     @torch.jit.export
     def lstm_sequence(self, seq: PackedSequence) -> Tuple[Tensor, Tensor]:
-        # if not torch.jit.is_scripting():
-        #    self.bilstm.flatten_parameters()
+        if not torch.jit.is_scripting():
+            self.bilstm.flatten_parameters()
         ret, _ = self.bilstm(seq)
         return nn.utils.rnn.pad_packed_sequence(ret, batch_first=True)