Add decoder block

test/test_models.py

@@ -1,5 +1,5 @@
 import torch
-from torchaudio.models import Wav2Letter, MelResNet, UpsampleNetwork, WaveRNN
+from torchaudio.models import Wav2Letter, MelResNet, UpsampleNetwork, WaveRNN, _Encoder
 
 from . import common_utils
 
@@ -115,3 +115,23 @@ def test_waveform(self):
         out = model(x, mels)
 
         assert out.size() == (n_batch, 1, hop_length * (n_time - kernel_size + 1), n_classes)
+
+
+class TestEncoder(common_utils.TorchaudioTestCase):
+    def test_output(self):
+        """Validate the output dimensions of a _Encoder block.
+        """
+
+        n_encoder_convolutions = 3
+        n_encoder_embedding = 512
+        n_encoder_kernel_size = 5
+        n_batch = 32
+        n_seq = 64
+
+        model = _Encoder(n_encoder_convolutions, n_encoder_embedding, n_encoder_kernel_size)
+
+        x = torch.rand(n_batch, n_encoder_embedding, n_seq)
+        input_length = [n_seq for i in range(n_batch)]
+        out = model(x, input_length)
+
+        assert out.size() == (n_batch, n_seq, n_encoder_embedding)

torchaudio/models/__init__.py

@@ -1,2 +1,3 @@
 from .wav2letter import *
 from .wavernn import *
+from ._tacotron import *