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.

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