diff --git a/paddle/fluid/train/custom_trainer/feed/scripts/create_programs.py b/paddle/fluid/train/custom_trainer/feed/scripts/create_programs.py
index c9b9ca94bc9d20283a2f5accf7b1470e2f73da30..e2b1cb7caa7d9a2589a448cd046812ad34805a6c 100644
--- a/paddle/fluid/train/custom_trainer/feed/scripts/create_programs.py
+++ b/paddle/fluid/train/custom_trainer/feed/scripts/create_programs.py
@@ -114,12 +114,14 @@ class ModelBuilder:
             with open(os.path.join(self._save_path, name), 'w') as f:
                 f.write(program.desc.serialize_to_string())
 
+        params = filter(fluid.io.is_parameter, main_program.list_vars())
+
         model_desc_path = os.path.join(self._save_path, 'model.yaml')
         model_desc = {
             'inputs': [{"name": var.name, "shape": var.shape} for var in inputs],
             'outputs': [{"name": var.name, "shape": var.shape} for var in outputs],
             'labels': [{"name": var.name, "shape": var.shape} for var in labels],
-            'vars': [{"name": var.name, "shape": var.shape} for var in main_program.list_vars() if fluid.io.is_parameter(var)],
+            'vars': [{"name": var.name, "shape": var.shape} for var in params],
             'loss': loss.name,
         }