bahdanau attn dropout was ignored

brainsqueeze · Jul 15, 2022 · 7b4fff7 · 7b4fff7
1 parent 85e44b7
commit 7b4fff7
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 
 setup(
     name="text2vec",
-    version="2.0.2",
+    version="2.0.3",
     description="Building blocks for text vectorization and embedding",
     author="Dave Hollander",
     author_url="https://github.com/brainsqueeze",

diff --git a/text2vec/models/transformer.py b/text2vec/models/transformer.py
@@ -71,7 +71,7 @@ def call(self, x, mask, training: bool = False):
             x = self.h_drop(ffn(x), training=training) + x
             x = self.layer_norm(x)
 
-        x, context = self.attention(x)
+        x, context = self.attention(x, training=training)
         return x, context
 
 
@@ -124,7 +124,7 @@ def call(self, x_enc, x_dec, dec_mask, context, attention: BahdanauAttention, tr
             ), training=training) + x_dec
             x_dec = self.layer_norm(x_dec)
 
-            x_dec, cross_context = attention(encoded=x_enc, decoded=x_dec)
+            x_dec, cross_context = attention(encoded=x_enc, decoded=x_dec, training=training)
             x_dec = self.h_drop(self.projection(x_dec, projection_vector=cross_context), training=training) + x_dec
 
             x_dec = self.layer_norm(x_dec)