Skip to content

Commit 355bc79

Browse files
author
Ruo-Ping Dong
authored
Fix gpu pytests (#5019)
* Move tensors to cpu before converting it to numpy
1 parent 6c7719a commit 355bc79

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

ml-agents/mlagents/trainers/optimizer/torch_optimizer.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ def _evaluate_by_sequence(
8888
# For the first sequence, the initial memory should be the one at the
8989
# beginning of this trajectory.
9090
for _ in range(first_seq_len):
91-
all_next_memories.append(initial_memory.squeeze().detach().numpy())
91+
all_next_memories.append(ModelUtils.to_numpy(initial_memory.squeeze()))
9292

9393
init_values, _mem = self.critic.critic_pass(
9494
seq_obs, initial_memory, sequence_length=first_seq_len
@@ -105,7 +105,7 @@ def _evaluate_by_sequence(
105105
):
106106
seq_obs = []
107107
for _ in range(self.policy.sequence_length):
108-
all_next_memories.append(_mem.squeeze().detach().numpy())
108+
all_next_memories.append(ModelUtils.to_numpy(_mem.squeeze()))
109109
for _obs in tensor_obs:
110110
start = seq_num * self.policy.sequence_length - (
111111
self.policy.sequence_length - leftover

0 commit comments

Comments
 (0)