Update transformer

1 files changed, 4 insertions, 6 deletions

diff --git a/text_recognizer/networks/transformer/decoder_block.py b/text_recognizer/networks/transformer/decoder_block.py
index b80a078..2dc4ddf 100644
--- a/text_recognizer/networks/transformer/decoder_block.py
+++ b/text_recognizer/networks/transformer/decoder_block.py
@@ -5,6 +5,7 @@ from typing import Optional, Type
 from torch import Tensor, nn
 
 from text_recognizer.networks.transformer.attention import Attention
+from text_recognizer.networks.transformer.embeddings.rotary import RotaryEmbedding
 from text_recognizer.networks.transformer.ff import FeedForward
 
 
@@ -25,22 +26,19 @@ class DecoderBlock(nn.Module):
         self.cross_attn = cross_attn
         self.ln_ff = deepcopy(norm)
         self.ff = ff
-        self.has_pos_emb = self.attn.rotary_embedding is not None
 
     def forward(
         self,
         x: Tensor,
         context: Optional[Tensor] = None,
-        input_mask: Optional[Tensor] = None,
-        context_mask: Optional[Tensor] = None,
+        mask: Optional[Tensor] = None,
+        rotary_embedding: Optional[RotaryEmbedding] = None,
     ) -> Tensor:
         """Applies decoder block on input signals."""
-        x = x + self.attn(self.ln_attn(x), input_mask=input_mask)
+        x = x + self.attn(self.ln_attn(x), mask=mask, rotary_embedding=rotary_embedding)
         x = x + self.cross_attn(
             x=self.ln_cross_attn(x),
             context=context,
-            input_mask=input_mask,
-            context_mask=context_mask,
         )
         x = x + self.ff(self.ln_ff(x))
         return x