From ddc6feabd5657d322a8cb7a82c471be7e88b7d90 Mon Sep 17 00:00:00 2001
From: LielinJiang <50691816+LielinJiang@users.noreply.github.com>
Date: Mon, 31 Jul 2023 14:45:37 +0800
Subject: [PATCH] polish kl_div doc (#55781)

---
 python/paddle/nn/functional/loss.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/python/paddle/nn/functional/loss.py b/python/paddle/nn/functional/loss.py
index 9a258458f2e..655fc8d5b7c 100644
--- a/python/paddle/nn/functional/loss.py
+++ b/python/paddle/nn/functional/loss.py
@@ -1610,7 +1610,10 @@ def kl_div(input, label, reduction='mean', name=None):
             import paddle.nn.functional as F
 
             shape = (5, 20)
-            x = paddle.uniform(shape, min=-10, max=10).astype('float32')
+
+            # input(x) should be a distribution in the log space
+            x = F.log_softmax(paddle.randn(shape), axis=1).astype('float32')
+
             target = paddle.uniform(shape, min=-10, max=10).astype('float32')
 
             # 'batchmean' reduction, loss shape will be [], who is 0-D Tensor
-- 
GitLab