summaryrefslogtreecommitdiff
path: root/text_recognizer/models/transformer.py
diff options
context:
space:
mode:
authorGustaf Rydholm <gustaf.rydholm@gmail.com>2022-09-28 23:44:06 +0200
committerGustaf Rydholm <gustaf.rydholm@gmail.com>2022-09-28 23:44:06 +0200
commit1bc811673ffc155d102a0e0ef74eff05c2ab4c13 (patch)
treef38b62cc0c8d71c11f060c0c162be5979876d401 /text_recognizer/models/transformer.py
parent923e18913a6518d6278d1fc1843a01bacf955c60 (diff)
Fix decode bug
Diffstat (limited to 'text_recognizer/models/transformer.py')
-rw-r--r--text_recognizer/models/transformer.py4
1 files changed, 2 insertions, 2 deletions
diff --git a/text_recognizer/models/transformer.py b/text_recognizer/models/transformer.py
index b2e5d5f..f6f10a7 100644
--- a/text_recognizer/models/transformer.py
+++ b/text_recognizer/models/transformer.py
@@ -95,8 +95,8 @@ class LitTransformer(LitBase):
Tensor: A tensor of token indices of the predictions from the model.
"""
start_index = self.tokenizer.start_index
- end_index = self.tokenizer.start_index
- pad_index = self.tokenizer.start_index
+ end_index = self.tokenizer.end_index
+ pad_index = self.tokenizer.pad_index
bsz = x.shape[0]
# Encode image(s) to latent vectors.