From 20d435b1171faffe029fb2d77aa3c62378565218 Mon Sep 17 00:00:00 2001 From: seiriosPlus Date: Tue, 22 Sep 2020 16:58:00 +0800 Subject: [PATCH] add large scale size method for pybind --- paddle/fluid/pybind/communicator_py.cc | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/paddle/fluid/pybind/communicator_py.cc b/paddle/fluid/pybind/communicator_py.cc index 6ac37a85c28..07ba7061678 100644 --- a/paddle/fluid/pybind/communicator_py.cc +++ b/paddle/fluid/pybind/communicator_py.cc @@ -109,10 +109,15 @@ void BindLargeScaleKV(py::module* m) { auto* sparse_variable = self.Get(table_name); sparse_variable->Load(dir); }) - .def("save", [](LargeScaleKV& self, const std::string& table_name, - const std::string& dir) { + .def("save", + [](LargeScaleKV& self, const std::string& table_name, + const std::string& dir) { + auto* sparse_variable = self.Get(table_name); + sparse_variable->Save(dir); + }) + .def("size", [](LargeScaleKV& self, const std::string& table_name) { auto* sparse_variable = self.Get(table_name); - sparse_variable->Save(dir); + return sparse_variable->Size(); }); } } // namespace pybind -- GitLab