Fix Python IndexError of case8: paddle.static.nn.nce (#49989)

* add dimension check for nce * add unittest * fix incorrect type in test_nce
PaddlePaddle · Feb 2, 2023 · c8548af · c8548af
1 parent 24e395f
commit c8548af
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/python/paddle/fluid/tests/unittests/test_nce.py b/python/paddle/fluid/tests/unittests/test_nce.py
@@ -330,6 +330,13 @@ def test_errors(self):
                 TypeError, paddle.static.nn.nce, input4, label4, 5
             )
 
+            input5 = paddle.static.data(name='x', shape=[1], dtype='float32')
+            label5 = paddle.static.data(name='label', shape=[1], dtype='int64')
+
+            self.assertRaises(
+                ValueError, paddle.static.nn.nce, input5, label5, 1
+            )
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/python/paddle/static/nn/loss.py b/python/paddle/static/nn/loss.py
@@ -129,6 +129,11 @@ def nce(
     check_variable_and_dtype(input, 'input', ['float32', 'float64'], 'nce')
     check_variable_and_dtype(label, 'label', ['int64'], 'nce')
 
+    if input.ndim != 2:
+        raise ValueError(
+            f'The rank of `input` must be 2, but received {input.ndim}.'
+        )
+
     dim = input.shape[1]
     num_true_class = label.shape[1]
     w = helper.create_parameter(