fix bug for whole hbm mode (PaddlePaddle#110)

Co-authored-by: yangjunchao <yangjunchao@baidu.com>
danleifeng · Sep 13, 2022 · cc71f56 · cc71f56
1 parent 582d236
commit cc71f56
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/paddle/fluid/framework/fleet/ps_gpu_wrapper.cc b/paddle/fluid/framework/fleet/ps_gpu_wrapper.cc
@@ -1100,7 +1100,12 @@ void PSGPUWrapper::SparseTableToHbm() {
   add_key_to_local(vec_data);
   add_key_to_gputask(gpu_task);
   BuildPull(gpu_task);
-  PrepareGPUTask(gpu_task);
+  if (!multi_mf_dim_) {
+    PrepareGPUTask(gpu_task);
+  }
+  else {
+    divide_to_device(gpu_task);
+  }
   BuildGPUTask(gpu_task);
   current_task_ = gpu_task;
   hbm_sparse_table_initialized_ = true;