Torch sequence_mask now works for tensors on different devices (#7980)

2026-07-17 11:32:33 +08:00 · 2020-04-15 07:21:51 +02:00
parent c17404918c
commit 8e439688fc
1 changed files with 2 additions and 2 deletions
@@ -51,8 +51,8 @@ def sequence_mask(lengths, maxlen, dtype=None):
    if maxlen is None:
        maxlen = lengths.max()

-    mask = ~(torch.ones((len(lengths), maxlen)).cumsum(dim=1).t() > lengths). \
-        t()
+    mask = ~(torch.ones((len(lengths), maxlen)).to(
+        lengths.device).cumsum(dim=1).t() > lengths).t()
    mask.type(dtype or torch.bool)

    return mask