hiro/ray
1
0
Fork 0
mirror of https://github.com/vale981/ray synced 2025-03-06 02:21:39 -05:00
Code Issues Projects Releases Packages Wiki Activity Actions

[RLlib] Update alpha_zero_policy.py (#15042)

Browse source
This commit is contained in:
Yeachan-Heo 2021-05-04 20:20:24 +09:00 • committed by GitHub
parent 40fdedd3de
commit 0552f6e886
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
rllib/contrib/alpha_zero/core/alpha_zero_policy.py
View file

@ -138,9 +138,9 @@ class AlphaZeroPolicy(TorchPolicy):
grad_info = self.extra_grad_info(train_batch)
grad_info.update(grad_process_info)
grad_info.update({
"total_loss": loss_out.detach().numpy(),
"policy_loss": policy_loss.detach().numpy(),
"value_loss": value_loss.detach().numpy()
"total_loss": loss_out.detach().cpu().numpy(),
"policy_loss": policy_loss.detach().cpu().numpy(),
"value_loss": value_loss.detach().cpu().numpy()
})
return {LEARNER_STATS_KEY: grad_info}