diff options
| -rw-r--r-- | text_recognizer/networks/transformer/layers.py | 4 |
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) |