瀏覽代碼

Update mineru/model/utils/pytorchocr/modeling/heads/rec_ppformulanet_head.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Xiaomeng Zhao 4 周之前
父節點
當前提交
33d43a5afc
共有 1 個文件被更改,包括 2 次插入2 次删除
  1. 2 2
      mineru/model/utils/pytorchocr/modeling/heads/rec_ppformulanet_head.py

+ 2 - 2
mineru/model/utils/pytorchocr/modeling/heads/rec_ppformulanet_head.py

@@ -251,13 +251,13 @@ class AttentionMaskConverter:
         bsz, src_len = mask.shape
         tgt_len = tgt_len if tgt_len is not None else src_len
         expanded_mask = (
-            mask[:, None, None, :].expand([bsz, 1, tgt_len, src_len]).cast(dtype)
+            mask[:, None, None, :].expand([bsz, 1, tgt_len, src_len]).to(dtype)
         )
 
         inverted_mask = 1.0 - expanded_mask
 
         return inverted_mask.masked_fill_(
-            inverted_mask.cast(torch.bool), torch.finfo(dtype).min
+            inverted_mask.to(torch.bool), torch.finfo(dtype).min
         )
 
     def _expand_mask_export(self, mask, dtype, tgt_len=None):