diff --git a/rllib/agents/a3c/a3c_torch_policy.py b/rllib/agents/a3c/a3c_torch_policy.py index 3b3439aa1..8e29149ea 100644 --- a/rllib/agents/a3c/a3c_torch_policy.py +++ b/rllib/agents/a3c/a3c_torch_policy.py @@ -64,8 +64,11 @@ def apply_grad_clipping(policy, optimizer, loss): params = list( filter(lambda p: p.grad is not None, param_group["params"])) if params: - info["grad_gnorm"] = nn.utils.clip_grad_norm_( + grad_gnorm = nn.utils.clip_grad_norm_( params, policy.config["grad_clip"]) + if isinstance(grad_gnorm, torch.Tensor): + grad_gnorm = grad_gnorm.cpu().numpy() + info["grad_gnorm"] = grad_gnorm return info