From 1baeae6b414f71906bd1480d3ddc393ae878bd63 Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Sun, 2 May 2021 14:10:53 +0200
Subject: Working on attention

---
 text_recognizer/networks/transformer/attention.py | 7 +++++++
 1 file changed, 7 insertions(+)

(limited to 'text_recognizer/networks/transformer')

diff --git a/text_recognizer/networks/transformer/attention.py b/text_recognizer/networks/transformer/attention.py
index e1324af..8724691 100644
--- a/text_recognizer/networks/transformer/attention.py
+++ b/text_recognizer/networks/transformer/attention.py
@@ -58,6 +58,7 @@ class Attention(nn.Module):
         context_mask: Optional[Tensor],
         rotary_pos_emb: Optional[Tensor] = None,
     ) -> Tuple[Tensor, Tensor]:
+        b, n, _, device = x.shape, x.device
         q, k, v = self.qkv_fn(x)
         q, k = (
             self._apply_rotary_emb(q, k, rotary_pos_emb)
@@ -66,7 +67,13 @@ class Attention(nn.Module):
             k,
         )
 
+        input_mask = None
         if any(x is not None for x in (mask, context_mask)):
+            q_mask = (
+                mask
+                if mask is not None
+                else lambda: torch.ones((b, n), device=device).bool()
+            )
             pass
 
         # Compute the attention
-- 
cgit v1.2.3-70-g09d2