From 4a82a3b01750cc98a83cc99ada79bc3aa910a95c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BC=A0=E6=98=A5=E4=B9=94?=
 <83450930+Liyulingyue@users.noreply.github.com>
Date: Fri, 10 Mar 2023 16:30:31 +0800
Subject: [PATCH] [phi] Add register of yolo_loss (#51316)

* add register of yolo_loss_kernel.cc

* fix something

* Update paddle/phi/kernels/cpu/yolo_loss_kernel.cc
---
 .../framework/new_executor/interpreter/interpreter_util.cc    | 3 +--
 paddle/phi/kernels/cpu/yolo_loss_kernel.cc                    | 4 +++-
 2 files changed, 4 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 af08155a315..7a032dd0339 100644
--- a/paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc
+++ b/paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc
@@ -97,8 +97,7 @@ static std::set<std::string> OpsNeedSetOutputDtypeWhenRegisterPhiKernel = {
     "unique",
     "unique_consecutive_flattened_tensor",
     "unique_raw",
-    "viterbi_devode",
-    "yolo_loss"};
+    "viterbi_devode"};
 
 // These Ops can use InferMeta to infer the output dtype
 static std::set<std::string> OpsWithAvailablePhiInferMeta = {
diff --git a/paddle/phi/kernels/cpu/yolo_loss_kernel.cc b/paddle/phi/kernels/cpu/yolo_loss_kernel.cc
index b32d7ee5962..788f6bb4f0b 100644
--- a/paddle/phi/kernels/cpu/yolo_loss_kernel.cc
+++ b/paddle/phi/kernels/cpu/yolo_loss_kernel.cc
@@ -371,4 +371,6 @@ void YoloLossKernel(const Context& dev_ctx,
 }  // namespace phi
 
 PD_REGISTER_KERNEL(
-    yolo_loss, CPU, ALL_LAYOUT, phi::YoloLossKernel, float, double) {}
+    yolo_loss, CPU, ALL_LAYOUT, phi::YoloLossKernel, float, double) {
+  kernel->OutputAt(2).SetDataType(phi::DataType::INT32);
+}
-- 
GitLab