fix SelectedRows rows() method gpu runtime error

PaddlePaddle · Oct 16, 2017 · c65bdd9 · c65bdd9
1 parent 9706992
commit c65bdd9
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/paddle/pybind/pybind.cc b/paddle/pybind/pybind.cc
@@ -153,7 +153,17 @@ PYBIND11_PLUGIN(core) {
       .def("set_height", &SelectedRows::set_height)
       .def("height", &SelectedRows::height)
       .def("set_rows", &SelectedRows::set_rows)
-      .def("rows", &SelectedRows::rows, py::return_value_policy::reference);
+      .def("rows", [](SelectedRows &self) {
+#ifndef PADDLE_WITH_CUDA
+        return self.rows();
+#else
+         auto rows = self.rows();
+         std::vector<int64_t> new_rows;
+         new_rows.reserve(rows.size());
+         std::copy(rows.begin(), rows.end(), std::back_inserter(new_rows));
+         return new_rows;
+#endif
+      });
 
   py::class_<Variable>(m, "Variable", R"DOC(Variable Class.