From 543277caa5b87159845d11355a1bf839d20c5815 Mon Sep 17 00:00:00 2001
From: caojian05 <caojian5@huawei.com>
Date: Thu, 19 Aug 2021 14:21:42 +0800
Subject: [PATCH] fix dqn mean_reward error

---
 model_zoo/official/rl/dqn/src/agent.py | 4 ++--
 model_zoo/official/rl/dqn/train.py     | 3 +--
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/model_zoo/official/rl/dqn/src/agent.py b/model_zoo/official/rl/dqn/src/agent.py
index 16aeb479d75..78a3508ba0a 100644
--- a/model_zoo/official/rl/dqn/src/agent.py
+++ b/model_zoo/official/rl/dqn/src/agent.py
@@ -118,5 +118,5 @@ class Agent:
         q_next_numpy = q_next.asnumpy()
         tem_ = Tensor(np.max(q_next_numpy, axis=1).reshape(-1, 1))
         q_target = b_r + self.gamma * tem_
-        self.train_net(b_s, q_target, b_a)
-        
\ No newline at end of file
+        loss = self.train_net(b_s, q_target, b_a)
+        return loss
diff --git a/model_zoo/official/rl/dqn/train.py b/model_zoo/official/rl/dqn/train.py
index 40a1234028a..294dee2b712 100644
--- a/model_zoo/official/rl/dqn/train.py
+++ b/model_zoo/official/rl/dqn/train.py
@@ -89,7 +89,7 @@ if __name__ == "__main__":
             agent.store_transition(s, a, r, s_)
             ep_r += r
             if agent.memory_counter > cfg.memory_capacity:
-                agent.learn()
+                _ = agent.learn()
                 if done_:
                     print("episode", episode, "total_reward", round(ep_r, 2))
                     rewards.append(round(ep_r, 2))
@@ -105,4 +105,3 @@ if __name__ == "__main__":
     times_numpy = np.array(times)
 
     print(rewards_numpy.mean(), times_numpy.mean())
-    
\ No newline at end of file