diff --git a/rllib/utils/torch_ops.py b/rllib/utils/torch_ops.py
index 6eaa14c27..87ce51a4e 100644
--- 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