diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc index 04087cb241c9cd4975773e646bc0ef6e1287518f..be5effe35fb90df928e4ca282557b8ff86b50bec 100644 --- a/paddle/fluid/pybind/pybind.cc +++ b/paddle/fluid/pybind/pybind.cc @@ -1663,6 +1663,15 @@ All parameter, weight, gradient are variables in Paddle. m.def("is_compiled_with_mkldnn", IsCompiledWithMKLDNN); m.def("is_compiled_with_brpc", IsCompiledWithBrpc); m.def("is_compiled_with_dist", IsCompiledWithDIST); + m.def("_cuda_synchronize", [](const platform::XPUPlace &place) { + platform::DeviceContextPool::Instance().Get(place)->Wait(); + }); + m.def("_cuda_synchronize", [](const platform::CUDAPlace &place) { + platform::DeviceContextPool::Instance().Get(place)->Wait(); + }); + m.def("_cuda_synchronize", [](const platform::CPUPlace &place) { + platform::DeviceContextPool::Instance().Get(place)->Wait(); + }); m.def("get_float_stats", []() { std::vector> float_stats; diff --git a/python/paddle/fluid/__init__.py b/python/paddle/fluid/__init__.py index e8cc6ce99016075a950f13d9e23f2957c9686471..9e60044a520d78d114f6d25116436aa81b5990e4 100644 --- a/python/paddle/fluid/__init__.py +++ b/python/paddle/fluid/__init__.py @@ -90,6 +90,7 @@ from .io import save, load, load_program_state, set_program_state from .dygraph.checkpoint import save_dygraph, load_dygraph from .dygraph.varbase_patch_methods import monkey_patch_varbase from . import generator +from .core_avx import _cuda_synchronize Tensor = LoDTensor enable_imperative = enable_dygraph disable_imperative = disable_dygraph @@ -133,7 +134,8 @@ __all__ = framework.__all__ + executor.__all__ + \ 'install_check', 'save', 'load', - 'VarBase' + 'VarBase', + '_cuda_synchronize' ]