From 627a6b8bacc5f4898c1c3c9018fd8e70ef95d8dc Mon Sep 17 00:00:00 2001 From: tangwei12 Date: Thu, 6 Dec 2018 17:14:59 +0800 Subject: [PATCH] add prefetch in nce --- paddle/fluid/operators/nce_op.h | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/paddle/fluid/operators/nce_op.h b/paddle/fluid/operators/nce_op.h index 8f82f77f5..7397d9f47 100644 --- a/paddle/fluid/operators/nce_op.h +++ b/paddle/fluid/operators/nce_op.h @@ -26,6 +26,10 @@ limitations under the License. */ #include "paddle/fluid/operators/math/sampler.h" #include "unsupported/Eigen/CXX11/Tensor" +#ifdef PADDLE_WITH_DISTRIBUTE +#include "paddle/fluid/operators/distributed/parameter_prefetch.h" +#endif + namespace paddle { namespace operators { @@ -166,8 +170,8 @@ class NCEKernel : public framework::OpKernel { auto height_sections = context.Attr>("height_sections"); auto table_names = context.Attr>("table_names"); - framework::Variable *ids = local_scope.Var("Ids"); - framework::Variable *weight = local_scope.Var("Weight"); + local_scope.Var("Ids"); + local_scope.Var("Weight"); #ifdef PADDLE_WITH_DISTRIBUTE operators::distributed::prefetch("Ids", "Weight", table_names, epmap, -- GitLab