Adding device

SforAiDl · Oct 21, 2020 · 4b11c16 · 4b11c16
1 parent f5a189d
commit 4b11c16
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/genrl/agents/modelbased/cem/cem.py b/genrl/agents/modelbased/cem/cem.py
@@ -60,7 +60,7 @@ def _create_model(self):
             "V",
             discrete,
             action_lim,
-        )
+        ).to(self.device)
         self.optim = torch.optim.Adam(self.agent.parameters(), lr=self.lr_policy)
 
     def plan(self):
@@ -136,7 +136,7 @@ def update_params(self):
         elite_states, elite_actions = self.select_elites(
             batch_states, batch_actions, batch_rewards
         )
-        action_probs = self.agent.forward(elite_states.float())
+        action_probs = self.agent.forward(elite_states.float().to(self.device))
         loss = F.cross_entropy(
             action_probs.view(-1, self.action_dim),
             elite_actions.long().view(-1),