From e72838c03d87c8eb10f7cb8df8804496b14121ab Mon Sep 17 00:00:00 2001
From: Julius Frost <33183774+juliusfrost@users.noreply.github.com>
Date: Wed, 9 Sep 2020 21:52:55 -0400
Subject: [PATCH] [RLLib] Add missing .to() for MARWIL on PyTorch (#10685)

There was a missing .to() that caused a device mismatch error on PyTorch with MARWIL.
---
 rllib/agents/marwil/marwil_torch_policy.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/rllib/agents/marwil/marwil_torch_policy.py b/rllib/agents/marwil/marwil_torch_policy.py
index fa2452d92..f10e71c2d 100644
--- a/rllib/agents/marwil/marwil_torch_policy.py
+++ b/rllib/agents/marwil/marwil_torch_policy.py
@@ -70,7 +70,7 @@ def stats(policy, train_batch):
 def setup_mixins(policy, obs_space, action_space, config):
     # Create a var.
     policy.ma_adv_norm = torch.tensor(
-        [100.0], dtype=torch.float32, requires_grad=False)
+        [100.0], dtype=torch.float32, requires_grad=False).to(policy.device)
     # Setup Value branch of our NN.
     ValueNetworkMixin.__init__(policy)