diff options
| -rw-r--r-- | text_recognizer/networks/perceiver/attention.py | 10 |
1 files changed, 1 insertions, 9 deletions
diff --git a/text_recognizer/networks/perceiver/attention.py b/text_recognizer/networks/perceiver/attention.py index 0ee51b1..19e3e17 100644 --- a/text_recognizer/networks/perceiver/attention.py +++ b/text_recognizer/networks/perceiver/attention.py @@ -25,9 +25,7 @@ class Attention(nn.Module): self.to_kv = nn.Linear(context_dim, 2 * inner_dim, bias=False) self.to_out = nn.Linear(inner_dim, query_dim, bias=False) - def forward( - self, x: Tensor, context: Optional[Tensor] = None, mask=Optional[Tensor] - ) -> Tensor: + def forward(self, x: Tensor, context: Optional[Tensor] = None) -> Tensor: h = self.heads q = self.to_q(x) context = context if context is not None else x @@ -36,12 +34,6 @@ class Attention(nn.Module): q, v, k = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v)) sim = einsum("b i d, b j d -> b i j", q, k) * self.scale - # if mask is not None: - # mask = rearrange(mask, "b ... -> b (...)") - # max_neg_val = -torch.finfo(sim.dtype).max - # mask = repeat(mask, "b j -> (b h) () j", h=h) - # sim.masked_fill_(~mask, max_neg_val) - attn = sim.softmax(dim=-1) out = einsum("b i j, b j d -> b i d", attn, v) out = rearrange(out, "(b h) n d -> b n (h d)", h=h) |