Fix output from attn modules

1 files changed, 2 insertions, 2 deletions

diff --git a/text_recognizer/networks/transformer/layers.py b/text_recognizer/networks/transformer/layers.py
index 941c141..f740244 100644
--- a/text_recognizer/networks/transformer/layers.py
+++ b/text_recognizer/networks/transformer/layers.py
@@ -91,9 +91,9 @@ class AttentionLayers(nn.Module):
                 x = norm(x)
 
             if layer_type == "a":
-                out, _ = block(x=x, mask=mask)
+                out = block(x=x, mask=mask)
             elif layer_type == "c":
-                out, _ = block(x, context=context, mask=mask, context_mask=context_mask)
+                out = block(x, context=context, mask=mask, context_mask=context_mask)
             elif layer_type == "f":
                 out = block(x)