diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc index 04087cb241c9cd4975773e646bc0ef6e1287518f..ae0b50461a113dd90a0b457b5a26030394657f20 100644 --- a/paddle/fluid/pybind/pybind.cc +++ b/paddle/fluid/pybind/pybind.cc @@ -1663,6 +1663,9 @@ 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::CUDAPlace &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..881fb86266fced5ea7d9f64a3507e58047c7f477 100644 --- a/python/paddle/fluid/__init__.py +++ b/python/paddle/fluid/__init__.py @@ -90,6 +90,8 @@ 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 import _cuda_synchronize + Tensor = LoDTensor enable_imperative = enable_dygraph disable_imperative = disable_dygraph @@ -133,7 +135,8 @@ __all__ = framework.__all__ + executor.__all__ + \ 'install_check', 'save', 'load', - 'VarBase' + 'VarBase', + '_cuda_synchronize' ] diff --git a/python/paddle/fluid/core.py b/python/paddle/fluid/core.py index 2e3bb6b00218a47467c59a2a88ea45ec80c32419..9a14c4cdf14a40bf5c23476b7bf21d8838b91251 100644 --- a/python/paddle/fluid/core.py +++ b/python/paddle/fluid/core.py @@ -263,6 +263,7 @@ if avx_supported(): from .core_avx import _save_dygraph_dict from .core_avx import _load_dygraph_dict from .core_avx import _create_loaded_parameter + from .core_avx import _cuda_synchronize if sys.platform != 'win32': from .core_avx import _set_process_pids from .core_avx import _erase_process_pids @@ -307,6 +308,7 @@ if load_noavx: from .core_noavx import _save_dygraph_dict from .core_noavx import _load_dygraph_dict from .core_noavx import _create_loaded_parameter + from .core_noavx import _cuda_synchronize if sys.platform != 'win32': from .core_noavx import _set_process_pids from .core_noavx import _erase_process_pids