summaryrefslogtreecommitdiff
path: root/text_recognizer/networks
diff options
context:
space:
mode:
Diffstat (limited to 'text_recognizer/networks')
-rw-r--r--text_recognizer/networks/transformer/layers.py4
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)