Decoder module working

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