Update rec_nrtr_loss.py

bupt906 · Oct 27, 2021 · 81f0a83 · 81f0a83
1 parent da02c71
commit 81f0a83
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ppocr/losses/rec_nrtr_loss.py b/ppocr/losses/rec_nrtr_loss.py
@@ -22,7 +22,7 @@ def forward(self, pred, batch):
             log_prb = F.log_softmax(pred, axis=1)
             non_pad_mask = paddle.not_equal(
                 tgt, paddle.zeros(
-                    tgt.shape, dtype='int32'))
+                    tgt.shape, dtype=tgt.dtype))
             loss = -(one_hot * log_prb).sum(axis=1)
             loss = loss.masked_select(non_pad_mask).mean()
         else: