diff --git a/paddle/fluid/framework/fleet/heter_ps/.CMakeLists.txt.swp b/paddle/fluid/framework/fleet/heter_ps/.CMakeLists.txt.swp
deleted file mode 100644
index 7d3f69e7424d33094dfdd9a2da0d3110a4895c8d..0000000000000000000000000000000000000000
Binary files a/paddle/fluid/framework/fleet/heter_ps/.CMakeLists.txt.swp and /dev/null differ
diff --git a/paddle/fluid/framework/fleet/heter_ps/graph_gpu_wrapper.cu b/paddle/fluid/framework/fleet/heter_ps/graph_gpu_wrapper.cu
index e99a0f4fe11c173f40f2764f49a5cbab695ea476..b0899b4a7f5b3f810a9a1f15f213f1282c996467 100644
--- a/paddle/fluid/framework/fleet/heter_ps/graph_gpu_wrapper.cu
+++ b/paddle/fluid/framework/fleet/heter_ps/graph_gpu_wrapper.cu
@@ -288,18 +288,18 @@ std::vector<int64_t> GraphGpuWrapper::graph_neighbor_sample(
   }
   /* VLOG(0) << "cumsum " << cumsum; */
 
-  std::vector<int64_t> res;
-  res.resize(cumsum * 2);
-  int count = 0;
+  std::vector<int64_t> cpu_key, res;
+  cpu_key.resize(key.size() * sample_size);
+
+  cudaMemcpy(cpu_key.data(), neighbor_sample_res.val,
+             key.size() * sample_size * sizeof(int64_t),
+             cudaMemcpyDeviceToHost);
   for (int i = 0; i < key.size(); i++) {
     for (int j = 0; j < actual_sample_size[i]; j++) {
-      res[count] = key[i];
-      count += 1;
+      res.push_back(key[i]);
+      res.push_back(cpu_key[i * sample_size + j]);
     }
   }
-
-  cudaMemcpy(res.data() + cumsum, neighbor_sample_res.val,
-             cumsum * sizeof(int64_t), cudaMemcpyDeviceToHost);
   /* for(int i = 0;i < res.size();i ++) { */
   /*     VLOG(0) << i << " " << res[i]; */
   /* } */