diff --git a/python/paddle/fluid/transpiler/distribute_transpiler.py b/python/paddle/fluid/transpiler/distribute_transpiler.py
index 27992df462ffd00ddf445538cc508b4232712481..c7ab300e0f0704ad16c15fce6fa3703587ff7c9e 100644
--- a/python/paddle/fluid/transpiler/distribute_transpiler.py
+++ b/python/paddle/fluid/transpiler/distribute_transpiler.py
@@ -177,6 +177,7 @@ class DistributeTranspiler:
                         dtype=table_grad_var.dtype)
                     for index in range(len(self.pserver_endpoints))
                 ]
+        return param_list, grad_list
 
     def _init_splited_vars(self, slice_var_up):
         # update these mappings for further transpile:
@@ -199,8 +200,8 @@ class DistributeTranspiler:
                 grad_list.append(g)
                 param_grad_set.add(g.name)
 
-        self._update_dist_lookup_table_vars(param_list, grad_list,
-                                            self.params_grads)
+        param_list, grad_list = self._update_dist_lookup_table_vars(
+            param_list, grad_list, self.params_grads)
 
         if slice_var_up:
             # when we slice var up into blocks, we will slice the var according to