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

rllib/utils/torch_ops.py

@@ -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