diff --git a/paddle/fluid/operators/custom_device_common_op_registry.cc b/paddle/fluid/operators/custom_device_common_op_registry.cc index 704d85acf13621fd33b852176c59eddd2b72f4ce..69625c03dbabec5c8b266ed9d6fb832d8e0236e9 100644 --- a/paddle/fluid/operators/custom_device_common_op_registry.cc +++ b/paddle/fluid/operators/custom_device_common_op_registry.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. */ #include "paddle/fluid/operators/custom_device_common_op_registry.h" +#include "paddle/fluid/operators/load_combine_op.h" #include "paddle/fluid/operators/run_program_op.h" #include "paddle/fluid/operators/save_combine_op.h" #include "paddle/phi/backends/device_manager.h" @@ -52,6 +53,19 @@ void RegisterCustomDeviceCommonKernel(const std::string& dev_type) { SaveCombineOpKernel, paddle::operators :: SaveCombineOpKernel); + REGISTER_OP_CUSTOM_DEVICE_KERNEL( + load_combine, + device_type, + paddle::operators:: + LoadCombineOpKernel, + paddle::operators:: + LoadCombineOpKernel, + paddle::operators:: + LoadCombineOpKernel, + paddle::operators:: + LoadCombineOpKernel, + paddle::operators:: + LoadCombineOpKernel); } } // namespace operators