From 8e439688fcc9acef272d988d52535ae2deb9983e Mon Sep 17 00:00:00 2001
From: Jan Blumenkamp <jan.blumenkamp@gmx.de>
Date: Wed, 15 Apr 2020 07:21:51 +0200
Subject: [PATCH] Torch sequence_mask now works for tensors on different
 devices (#7980)

---
 rllib/utils/torch_ops.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

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