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 d6e193d2878..defd34d6fe2 100644
--- a/model_zoo/official/rl/dqn/train.py
+++ b/model_zoo/official/rl/dqn/train.py
@@ -88,7 +88,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))