diff --git a/ppocr/modeling/heads/rec_sar_head.py b/ppocr/modeling/heads/rec_sar_head.py index 647f58200f83e3f2c23030872c606f299c58f7b7..7107788d9ef3b49ac6d4dcd4a8133a9603ada19b 100644 --- a/ppocr/modeling/heads/rec_sar_head.py +++ b/ppocr/modeling/heads/rec_sar_head.py @@ -235,7 +235,8 @@ class ParallelSARDecoder(BaseDecoder): # cal mask of attention weight for i, valid_ratio in enumerate(valid_ratios): valid_width = min(w, math.ceil(w * valid_ratio)) - attn_weight[i, :, :, valid_width:, :] = float('-inf') + if valid_width < w: + attn_weight[i, :, :, valid_width:, :] = float('-inf') attn_weight = paddle.reshape(attn_weight, [bsz, T, -1]) attn_weight = F.softmax(attn_weight, axis=-1)