From 545e20f8fbd7edaeac969c48c6f399a20ae7e7cd Mon Sep 17 00:00:00 2001 From: Sanbu <96160062+sanbuphy@users.noreply.github.com> Date: Sun, 19 Mar 2023 10:07:42 +0800 Subject: [PATCH] [phi] Add output defs for argsort kernel (#51407) * Add output defs for argsort kernel * Update argsort_kernel.cc * Update argsort_kernel.cu * Update argsort_kernel.cc --- .../framework/new_executor/interpreter/interpreter_util.cc | 1 - paddle/phi/kernels/cpu/argsort_kernel.cc | 1 + paddle/phi/kernels/gpu/argsort_kernel.cu | 4 +++- paddle/phi/kernels/xpu/argsort_kernel.cc | 4 +++- 4 files changed, 7 insertions(+), 3 deletions(-) diff --git a/paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc b/paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc index e4719cbcf0e..c434a3477e0 100644 --- a/paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc +++ b/paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc @@ -54,7 +54,6 @@ static std::set OpsNeedSetOutputDtypeWhenRegisterPhiKernel = { "adam", "adamw", "any_raw", - "arg_sort", "clip_by_norm", "eig_grad", "eigh", diff --git a/paddle/phi/kernels/cpu/argsort_kernel.cc b/paddle/phi/kernels/cpu/argsort_kernel.cc index 97f8fb67ed1..237fe6dc615 100644 --- a/paddle/phi/kernels/cpu/argsort_kernel.cc +++ b/paddle/phi/kernels/cpu/argsort_kernel.cc @@ -154,4 +154,5 @@ void ArgsortKernel(const Context& dev_ctx, PD_REGISTER_KERNEL( argsort, CPU, ALL_LAYOUT, phi::ArgsortKernel, float, double, int, int64_t) { + kernel->OutputAt(1).SetDataType(phi::DataType::INT64); } diff --git a/paddle/phi/kernels/gpu/argsort_kernel.cu b/paddle/phi/kernels/gpu/argsort_kernel.cu index 13455a7639c..0ad699cdc64 100644 --- a/paddle/phi/kernels/gpu/argsort_kernel.cu +++ b/paddle/phi/kernels/gpu/argsort_kernel.cu @@ -519,4 +519,6 @@ PD_REGISTER_KERNEL(argsort, double, int, int64_t, - phi::dtype::float16) {} + phi::dtype::float16) { + kernel->OutputAt(1).SetDataType(phi::DataType::INT64); +} diff --git a/paddle/phi/kernels/xpu/argsort_kernel.cc b/paddle/phi/kernels/xpu/argsort_kernel.cc index 3324af07d0d..e1875b8f52c 100644 --- a/paddle/phi/kernels/xpu/argsort_kernel.cc +++ b/paddle/phi/kernels/xpu/argsort_kernel.cc @@ -248,4 +248,6 @@ PD_REGISTER_KERNEL(argsort, float, int, int64_t, - phi::dtype::float16) {} + phi::dtype::float16) { + kernel->OutputAt(1).SetDataType(phi::DataType::INT64); +} -- GitLab