fix unittest coverage

PaddlePaddle · Nov 17, 2021 · db2ffd4 · db2ffd4 · paddle-bot-old · Nov 17, 2021
1 parent bbcd325
commit db2ffd4
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/python/paddle/fluid/tests/unittests/test_kldiv_loss_op.py b/python/paddle/fluid/tests/unittests/test_kldiv_loss_op.py
@@ -112,7 +112,9 @@ def test_kl_loss_static_api(self):
         input = paddle.fluid.data(name='input', shape=[5, 20])
         label = paddle.fluid.data(name='label', shape=[5, 20])
 
-        pred_loss = paddle.nn.functional.kl_div(input, label)
+        paddle.nn.functional.kl_div(input, label)
+        paddle.nn.functional.kl_div(input, label, 'sum')
+        paddle.nn.functional.kl_div(input, label, 'batchmean')
 
 
 class TestKLDivLossTypePromotion(unittest.TestCase):

diff --git a/python/paddle/nn/functional/loss.py b/python/paddle/nn/functional/loss.py
@@ -935,9 +935,8 @@ def kl_div(input, label, reduction='mean', name=None):
     elif reduction == 'sum':
         loss = paddle.sum(loss)
     elif reduction == 'batchmean':
-        if len(paddle.shape(input)) > 0:
-            batch_size = paddle.shape(input)[0]
-            loss = paddle.sum(out) / batch_size
+        batch_size = paddle.shape(input)[0]
+        loss = paddle.sum(loss) / batch_size
     return loss