summaryrefslogtreecommitdiff
path: root/text_recognizer/models/transformer.py
diff options
context:
space:
mode:
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.