Polish cmakelists

PaddlePaddle · Jul 4, 2022 · 5a55f92 · 5a55f92
1 parent c1073aa
commit 5a55f92
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 18 deletions.
diff --git a/paddle/fluid/jit/CMakeLists.txt b/paddle/fluid/jit/CMakeLists.txt
@@ -26,7 +26,8 @@ cc_library(
 cc_library(
   jit_layer
   SRCS layer.cc
-  DEPS jit_compilation_unit)
+  DEPS jit_serializer jit_function_utils jit_serializer_utils
+       jit_compilation_unit jit_function_schema)
 
 if(WITH_TESTING
    AND NOT WIN32
@@ -45,12 +46,7 @@ if(WITH_TESTING
       feed_op
       fetch_op
       scale_op
-      jit_serializer
-      jit_layer
-      jit_function_utils
-      jit_function_schema
-      jit_compilation_unit
-      jit_serializer_utils)
+      jit_layer)
   cc_test(
     layer_test
     SRCS layer_test.cc

diff --git a/paddle/fluid/jit/executor_function.h b/paddle/fluid/jit/executor_function.h
@@ -56,6 +56,8 @@ class ExecutorFunction : public BaseFunction {
     return res;
   }
 
+  const std::shared_ptr<FunctionInfo> &Info() const { return info_; }
+
  private:
   std::shared_ptr<FunctionInfo> info_;
   framework::Scope scope_;

diff --git a/paddle/fluid/pybind/CMakeLists.txt b/paddle/fluid/pybind/CMakeLists.txt
@@ -39,11 +39,7 @@ set(PYBIND_DEPS
     phi_utils
     tcp_store
     new_profiler
-    jit_serializer
-    jit_layer
-    jit_function_utils
-    jit_serializer_utils
-    jit_function_schema)
+    jit_layer)
 
 if(WITH_PSCORE)
   set(PYBIND_DEPS ${PYBIND_DEPS} ps_service)

diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
@@ -1757,13 +1757,18 @@ PYBIND11_MODULE(core_noavx, m) {
                var_inputs.emplace_back(tensor->Var());
              }
              auto var_outputs = self(var_inputs);
+
              std::vector<std::shared_ptr<imperative::VarBase>> tensor_outputs;
-             for (auto &var : var_outputs) {
-               imperative::VariableWrapper var_wrapper("x", var);
-               auto tmp =
+             auto output_names = self.Info()->OutputArgNames();
+             for (size_t i = 0; i < var_outputs.size(); ++i) {
+               auto var = var_outputs[i];
+               std::string name = output_names[i];
+               imperative::VariableWrapper var_wrapper(name, var);
+               auto shared_wrapper =
                    std::make_shared<imperative::VariableWrapper>(var_wrapper);
-               auto tmp2 = std::make_shared<imperative::VarBase>(tmp);
-               tensor_outputs.emplace_back(tmp2);
+               auto shared_varbase =
+                   std::make_shared<imperative::VarBase>(shared_wrapper);
+               tensor_outputs.emplace_back(shared_varbase);
              }
              return tensor_outputs;
            });

diff --git a/python/paddle/fluid/tests/unittests/test_jit_layer.py b/python/paddle/fluid/tests/unittests/test_jit_layer.py
@@ -71,7 +71,6 @@ def test_multi_load(self):
         jit_layer = JitLayer(model_path, place)
         forward_out2 = jit_layer.forward(x)
         infer_out2 = jit_layer.infer(x)
-        # import pdb; pdb.set_trace()
         self.assertEqual(np.allclose(forward_out1, forward_out2[0]), True)
         self.assertEqual(np.allclose(infer_out1, infer_out2[0]), True)