diff options
author | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2021-05-13 23:02:20 +0200 |
---|---|---|
committer | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2021-05-13 23:02:20 +0200 |
commit | 8c7768e8d321efec558e12bff9b89b2de615d541 (patch) | |
tree | 67f5928c5584e8826c01834d06d34cd7e60546ba /text_recognizer/networks/transformer/transformer.py | |
parent | c9c60678673e19ad3367339eb8e7a093e5a98474 (diff) |
Decoder module working
Diffstat (limited to 'text_recognizer/networks/transformer/transformer.py')
-rw-r--r-- | text_recognizer/networks/transformer/transformer.py | 4 |
1 files changed, 2 insertions, 2 deletions
diff --git a/text_recognizer/networks/transformer/transformer.py b/text_recognizer/networks/transformer/transformer.py index 36f86ac..60ab1ce 100644 --- a/text_recognizer/networks/transformer/transformer.py +++ b/text_recognizer/networks/transformer/transformer.py @@ -1,5 +1,5 @@ """Transformer wrapper.""" -from typing import Optional, Type +from typing import Any, Optional, Type from torch import nn, Tensor @@ -58,4 +58,4 @@ class Transformer(nn.Module): x = self.project_emb(x) x = self.attn_layers(x, mask=mask, **kwargs) out = self.logits(x) if not return_embeddings else x - return x + return out |