fix scope in nce and prefetch

paddle/fluid/operators/distributed/parameter_prefetch.cc

@@ -102,8 +102,7 @@ static void MergeMultipleVarsIntoOneBySection(
     const std::string& out_name, const std::vector<std::string>& out_var_names,
     const std::vector<int>& height_section,
     const std::vector<std::vector<int64_t>>& splited_ids,
-    const framework::ExecutionContext& context,
-    const framework::Scope& actual_scope, framework::Scope* scope,
+    const framework::ExecutionContext& context, framework::Scope* scope,
     platform::DeviceContext* actual_ctx) {
   PADDLE_ENFORCE_EQ(out_var_names.size(), height_section.size(), "");
 
@@ -115,9 +114,9 @@ static void MergeMultipleVarsIntoOneBySection(
     id_to_offset[ids_vector[i]].push_back(i);
   }
 
-  auto& id_tensor = actual_scope.FindVar(id_name)->Get<framework::LoDTensor>();
+  auto& id_tensor = scope.FindVar(id_name)->Get<framework::LoDTensor>();
   auto* out_tensor =
-      actual_scope.FindVar(out_name)->GetMutable<framework::LoDTensor>();
+      scope.FindVar(out_name)->GetMutable<framework::LoDTensor>();
   auto* out_tensor_data = out_tensor->mutable_data<float>(id_tensor.place());
 
   bool is_on_cpu_place = true;
@@ -175,7 +174,7 @@ void prefetch(const std::string& id_name, const std::string& out_name,
               const std::vector<int>& height_sections,
               const framework::ExecutionContext& context,
               const framework::Scope& scope) {
-  auto& local_scope = context.scope().NewScope();
+  auto& local_scope = scope.NewScope();
 
   platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
   auto& cpu_ctx = *pool.Get(platform::CPUPlace());
@@ -247,8 +246,8 @@ void prefetch(const std::string& id_name, const std::string& out_name,
 
   MergeMultipleVarsIntoOneBySection(id_name, ids_vector, out_name,
                                     out_var_names, height_sections, splited_ids,
-                                    context, scope, &local_scope, &actual_ctx);
-  context.scope().DeleteScope(&local_scope);
+                                    context, &local_scope, &actual_ctx);
+  scope.DeleteScope(&local_scope);
 }
 
 };  // namespace distributed

paddle/fluid/operators/nce_op.h

@@ -170,7 +170,7 @@ class NCEKernel : public framework::OpKernel<T> {
       auto height_sections = context.Attr<std::vector<int>>("height_sections");
       auto table_names = context.Attr<std::vector<std::string>>("table_names");
 
-      auto *ids = local_scope.Var("Ids");
+      auto *ids = local_scope.Var("Ids@Local");
       auto *x_tensor = ids->GetMutable<framework::LoDTensor>();
       x_tensor->mutable_data<int64_t>(
           framework::make_ddim({static_cast<int64_t>(labels.size()), 1}),
@@ -179,12 +179,10 @@ class NCEKernel : public framework::OpKernel<T> {
       std::memcpy(x_tensor->data<int64_t>(), labels.data(),
                   labels.size() * sizeof(int64_t));
 
-      local_scope.Var("Weight@Local")
-          ->GetMutable<framework::LoDTensor>()
-          ->mutable_data<T>(context.GetPlace());
+      local_scope.Var("Weight@Local");
 
 #ifdef PADDLE_WITH_DISTRIBUTE
-      operators::distributed::prefetch("Ids", "Weight@Local", table_names,
+      operators::distributed::prefetch("Ids@Local", "Weight@Local", table_names,
                                        epmap, height_sections, context,
                                        &local_scope);
 #else
@@ -207,10 +205,7 @@ class NCEKernel : public framework::OpKernel<T> {
         sample_out_data[i] += result(0);
         sample_out_data[i] = (1. / (1. + exp(-sample_out_data[i])));
       }
-
-      if (context.scope().HasKid(&local_scope)) {
       context.scope().DeleteScope(&local_scope);
-      }
     } else {
       auto weight_mat =
           EigenMatrix<T>::From(*(context.Input<Tensor>("Weight")));