Stop token prediction - does train yet

tests/layers_tests.py

@@ -37,18 +37,23 @@ class DecoderTests(unittest.TestCase):
         dummy_memory = T.autograd.Variable(T.rand(4, 120, 32))
 
         print(layer)
-        output, alignment = layer(dummy_input, dummy_memory)
+        output, alignment, stop_output = layer(dummy_input, dummy_memory)
         print(output.shape)
+        print(" > Stop ", stop_output.shape)
+        
         assert output.shape[0] == 4
         assert output.shape[1] == 120 / 5
         assert output.shape[2] == 32 * 5
-
+        assert stop_output.shape[0] == 4
+        assert stop_output.shape[1] == 120 / 5
+        assert stop_output.shape[2] == 5
+        
 
 class EncoderTests(unittest.TestCase):
 
     def test_in_out(self):
         layer = Encoder(128)
-       dummy_input = T.autograd.Variable(T.rand(4, 8, 128))
+        dummy_input = T.autograd.Variable(T.rand(4, 8, 128))
 
         print(layer)
         output = layer(dummy_input)

tests/loader_tests.py

@@ -32,7 +32,7 @@ class TestDataset(unittest.TestCase):
                                   c.power
                                 )
 
-        dataloader = DataLoader(dataset, batch_size=c.batch_size,
+        dataloader = DataLoader(dataset, batch_size=2,
                                 shuffle=True, collate_fn=dataset.collate_fn,
                                 drop_last=True, num_workers=c.num_loader_workers)
 
@@ -43,7 +43,8 @@ class TestDataset(unittest.TestCase):
             text_lengths = data[1]
             linear_input = data[2]
             mel_input = data[3]
-            item_idx = data[4]
+            stop_targets = data[4]
+            item_idx = data[5]
 
             neg_values = text_input[text_input < 0]
             check_count = len(neg_values)
@@ -81,13 +82,16 @@ class TestDataset(unittest.TestCase):
             text_lengths = data[1]
             linear_input = data[2]
             mel_input = data[3]
-            item_idx = data[4]
+            stop_target = data[4]
+            item_idx = data[5]
 
             # check the last time step to be zero padded
             assert mel_input[0, -1].sum() == 0
             assert mel_input[0, -2].sum() != 0
             assert linear_input[0, -1].sum() == 0
             assert linear_input[0, -2].sum() != 0
+            assert stop_target[0, -1] == 1
+            assert stop_target.sum() == 1