Torch sequence_mask now works for tensors on different devices (ray-p…

…roject#7980)
aannadi · Apr 15, 2020 · 8e43968 · 8e43968
1 parent c174049
commit 8e43968
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/rllib/utils/torch_ops.py b/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