Fix typo in builders

idiap · Jul 20, 2021 · f22c137 · f22c137
1 parent 801ba89
commit f22c137
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/fast_transformers/builders/transformer_builders.py b/fast_transformers/builders/transformer_builders.py
@@ -357,7 +357,7 @@ def cross_query_dimensions(self):
         return self._cross_d_query or self.query_dimensions
 
     @cross_query_dimensions.setter
-    def cross_attention(self, n):
+    def cross_query_dimensions(self, n):
         self._cross_d_query = n
 
     @property

diff --git a/tests/test_builders.py b/tests/test_builders.py
@@ -106,6 +106,15 @@ def test_decoder_build(self):
         m = torch.rand(1, 13, 7*32)
         y = transformer(x, m)
 
+        t = TransformerDecoderBuilder.from_kwargs(
+            n_layers=1,
+            n_heads=4,
+            query_dimensions=32,
+            cross_n_heads=7,
+            cross_value_dimensions=32,
+            cross_query_dimensions=32
+        ).get()
+
     def test_recurrent_decoder(self):
         transformer = RecurrentDecoderBuilder().get()
         builder = RecurrentDecoderBuilder()