From f266318a017c3d4ca4a03ccc908bda05524c9c83 Mon Sep 17 00:00:00 2001
From: Tomasz Wrona <tomasz@wrona.me>
Date: Thu, 28 May 2020 21:06:27 +0200
Subject: [PATCH] [rllib] Do not store torch tensors when using grad clipping
 (#8509)

---
 rllib/agents/a3c/a3c_torch_policy.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

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