diff options
Diffstat (limited to 'text_recognizer/model/greedy_decoder.py')
| -rw-r--r-- | text_recognizer/model/greedy_decoder.py | 2 |
1 files changed, 1 insertions, 1 deletions
diff --git a/text_recognizer/model/greedy_decoder.py b/text_recognizer/model/greedy_decoder.py index 2c4c16e..8d55a02 100644 --- a/text_recognizer/model/greedy_decoder.py +++ b/text_recognizer/model/greedy_decoder.py @@ -34,7 +34,7 @@ class GreedyDecoder: for i in range(1, self.max_output_len): tokens = indecies[:, :i] # (B, Sy) logits = self.network.decode(tokens, img_features) # [ B, N, C ] - indecies_ = torch.argmax(logits, dim=2) # [ B, N ] + indecies_ = logits.argmax(dim=2) # [ B, N ] indecies[:, i] = indecies_[:, -1] # Early stopping of prediction loop if token is end or padding token. |