diff --git a/python/paddle/fluid/tests/unittests/test_auto_parallel_partitioner.py b/python/paddle/fluid/tests/unittests/test_auto_parallel_partitioner.py
index 0a4e8d3d9e2d4b8bf62e728aa45f2cd3c0a39af3..95063b85f7e5bca1935c9907c754df5ad4466543 100644
--- a/python/paddle/fluid/tests/unittests/test_auto_parallel_partitioner.py
+++ b/python/paddle/fluid/tests/unittests/test_auto_parallel_partitioner.py
@@ -234,9 +234,8 @@ def distributed_attr_check_for_dist_op(serial_main_prog, dist_main_prog,
 def distributed_attr_check_for_program(dist_main_prog, dist_context):
     have_dist_attr = True
     for block in dist_main_prog.blocks:
-        for tensor in block.vars.values():
-            var_dist_attr = dist_context.get_tensor_dist_attr_for_program(
-                tensor)
+        for var in block.vars.values():
+            var_dist_attr = dist_context.get_tensor_dist_attr_for_program(var)
             if var_dist_attr is None:
                 have_dist_attr = False