From 09f4e3b185de0074de94d4e9c26a3ca2fb67fa8f Mon Sep 17 00:00:00 2001 From: Tai An Date: Thu, 3 Sep 2026 02:09:26 -0700 Subject: [PATCH] fix(nrtr): let generate_square_subsequent_mask take the device it is given (#731) Transformer.forward_train calls tgt_mask = self.generate_square_subsequent_mask(tgt.shape[0], tgt.device) but the helper only accepts sz, so the training branch of Transformer.forward raises TypeError: Transformer.generate_square_subsequent_mask() takes 2 positional arguments but 3 were given The other two call sites (forward_test and forward_beam) pass one argument and are unaffected. Accept an optional device and forward it to the two tensor constructors. It defaults to None, so the one-argument call sites keep building the mask on the CPU exactly as before, and the call site that already asks for a device now gets one. --- .../pytorch/networks/heads/rec_nrtr_head.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/python/rapidocr/inference_engine/pytorch/networks/heads/rec_nrtr_head.py b/python/rapidocr/inference_engine/pytorch/networks/heads/rec_nrtr_head.py index cfcd073..33a69d4 100644 --- a/python/rapidocr/inference_engine/pytorch/networks/heads/rec_nrtr_head.py +++ b/python/rapidocr/inference_engine/pytorch/networks/heads/rec_nrtr_head.py @@ -355,13 +355,18 @@ class Transformer(nn.Module): torch.tensor(hyp_scores), ] - def generate_square_subsequent_mask(self, sz): + def generate_square_subsequent_mask(self, sz, device=None): """Generate a square mask for the sequence. The masked positions are filled with float('-inf'). Unmasked positions are filled with float(0.0). """ - mask = torch.zeros([sz, sz], dtype=torch.float32) + mask = torch.zeros([sz, sz], dtype=torch.float32, device=device) mask_inf = torch.triu( - torch.full(size=[sz, sz], fill_value=float("-Inf"), dtype=torch.float32), + torch.full( + size=[sz, sz], + fill_value=float("-Inf"), + dtype=torch.float32, + device=device, + ), diagonal=1, ) mask = mask + mask_inf