BatchedMcts validate_prior on CuArray

Whojo · Sep 24, 2022 · 90338dd · 90338dd
1 parent bf82e9e
commit 90338dd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/redesign/src/BatchedMcts.jl b/redesign/src/BatchedMcts.jl
@@ -313,7 +313,7 @@ function create_tree(mcts, envs)
     valid_actions = fill(false, mcts.device, (A, N, B))
     valid_actions[:, ROOT, :] = info.valid_actions
     policy_prior = zeros(Float32, mcts.device, (A, N, B))
-    policy_prior[:, ROOT, :] = validate_prior(info.policy_prior, info.valid_actions)
+    policy_prior[:, ROOT, :] = validate_prior(info.policy_prior, valid_actions[:, ROOT, :])
     value_prior = zeros(Float32, mcts.device, (N, B))
     value_prior[ROOT, :] = info.value_prior