Copy initial state of an RNN to a CPU before converting it to a NumPy…

… array (ray-project#8097)
aannadi · Apr 26, 2020 · b508166 · b508166
1 parent b506f87
commit b508166
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/rllib/policy/torch_policy.py b/rllib/policy/torch_policy.py
@@ -333,7 +333,9 @@ def num_state_tensors(self):
 
     @override(Policy)
     def get_initial_state(self):
-        return [s.numpy() for s in self.model.get_initial_state()]
+        return [
+            s.cpu().detach().numpy() for s in self.model.get_initial_state()
+        ]
 
     def extra_grad_process(self, optimizer, loss):
         """Called after each optimizer.zero_grad() + loss.backward() call.