diff --git a/python/paddle/fluid/incubate/fleet/parameter_server/ir/trainer_pass.py b/python/paddle/fluid/incubate/fleet/parameter_server/ir/trainer_pass.py
index 11fa70b70ba85236d507377de648fa63ec469be0..ff10c8ea097ad0ab3385018b89e3f0d9ec7f0408 100644
--- a/python/paddle/fluid/incubate/fleet/parameter_server/ir/trainer_pass.py
+++ b/python/paddle/fluid/incubate/fleet/parameter_server/ir/trainer_pass.py
@@ -771,7 +771,7 @@ def find_heter_ops(program, default_device="cpu"):
 
                 """
                 output_vars_no_grad = []
-                for key in pre_op.output_names:
+                for key in op.output_names:
                     for varname in op.output(key):
                         if varname == "@EMPTY@":
                             continue