bug fix for DeserializeSelectedRows. test=develop (PaddlePaddle#36520)

zmxdream · Nov 22, 2021 · f9f985a · f9f985a
1 parent 9ffb43b
commit f9f985a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/paddle/fluid/distributed/service/brpc_utils.cc b/paddle/fluid/distributed/service/brpc_utils.cc
@@ -273,8 +273,8 @@ void DeserializeSelectedRows(framework::Variable* var, const VarMsg& msg,
   auto* slr = var->GetMutable<framework::SelectedRows>();
   framework::Tensor* tensor = slr->mutable_value();
   slr->set_height(msg.slr_height());
-  std::vector<int64_t> tmp_rows(msg.slr_height());
-  memcpy(&tmp_rows[0], msg.data().data(), msg.slr_height() * sizeof(int64_t));
+  std::vector<int64_t> tmp_rows(msg.dims()[0]);
+  memcpy(tmp_rows.data(), msg.data().data(), msg.dims()[0] * sizeof(int64_t));
   slr->set_rows(tmp_rows);
   std::vector<int> vec_dim;
   for (auto& x : msg.dims()) {