diff --git a/python/paddle/fluid/tests/unittests/test_dist_fleet_ps4.py b/python/paddle/fluid/tests/unittests/test_dist_fleet_ps4.py
index 88ed835d7961a32b86aa8623b82e8cb8e7ab6d1f..f1176aea34ea88d821326597e34cd064fdbad26c 100644
--- a/python/paddle/fluid/tests/unittests/test_dist_fleet_ps4.py
+++ b/python/paddle/fluid/tests/unittests/test_dist_fleet_ps4.py
@@ -68,7 +68,7 @@ class TestPSPassWithBow(unittest.TestCase):
             name="query_ids", shape=[1], dtype="int64", lod_level=1)
         # embedding
         q_emb = fluid.contrib.layers.sparse_embedding(
-            input=pt,
+            input=q,
             size=[dict_dim, emb_dim],
             param_attr=fluid.ParamAttr(
                 initializer=fluid.initializer.Constant(value=0.01),
@@ -117,7 +117,7 @@ class TestPSPassWithBow(unittest.TestCase):
             name="neg_title_ids", shape=[1], dtype="int64", lod_level=1)
         # embedding
         nt_emb = fluid.contrib.layers.sparse_embedding(
-            input=pt,
+            input=nt,
             size=[dict_dim, emb_dim],
             param_attr=fluid.ParamAttr(
                 initializer=fluid.initializer.Constant(value=0.01),