Fix tests

CyberZHG · Sep 17, 2018 · 57bd1b2 · 57bd1b2
1 parent fc3017a
commit 57bd1b2
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/keras_self_attention/seq_pos_attention.py b/keras_self_attention/seq_pos_attention.py
diff --git a/tests/seq_weighted_attention/test_save_load.py b/tests/seq_weighted_attention/test_save_load.py
@@ -17,9 +17,16 @@ def _test_save_load(self, attention):
         lstm = keras.layers.Bidirectional(keras.layers.LSTM(units=7,
                                                             return_sequences=True),
                                           name='Bi-LSTM')(embd)
-        weighted, attention = attention(lstm)
-        dense = keras.layers.Dense(units=2, activation='softmax', name='Softmax')(weighted)
-        model = keras.models.Model(inputs=inputs, outputs=[dense, attention])
+        if attention.return_attention:
+            layer, weights = attention(lstm)
+        else:
+            layer = attention(lstm)
+        dense = keras.layers.Dense(units=2, activation='softmax', name='Softmax')(layer)
+        if attention.return_attention:
+            outputs = [dense, weights]
+        else:
+            outputs = dense
+        model = keras.models.Model(inputs=inputs, outputs=outputs)
         model.compile(
             optimizer='adam',
             loss={'Softmax': 'sparse_categorical_crossentropy'},
@@ -29,6 +36,10 @@ def _test_save_load(self, attention):
         model.save(model_path)
         model = keras.models.load_model(model_path, custom_objects=Attention.get_custom_objects())
         model.summary(line_length=100)
+        if attention.return_attention:
+            self.assertEqual(2, len(model.outputs))
+        else:
+            self.assertEqual(1, len(model.outputs))
 
     def test_default(self):
         self._test_save_load(Attention(name='Attention'))