diff options
| author | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2023-09-03 01:14:41 +0200 |
|---|---|---|
| committer | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2023-09-03 01:14:41 +0200 |
| commit | 53cfc21cffa4e877ad0959170b47b690d2fdb40f (patch) | |
| tree | 88eaabb75a5ca6fba3c43655ccb754a9f0c5d649 /text_recognizer/model/greedy_decoder.py | |
| parent | d020059f2f71fe7c25765dde9d535195c09ece01 (diff) | |
Update lit models
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. |