未验证 提交 4a82a3b0 编写于 作者: 张春乔 提交者: GitHub

[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
上级 fe264cfc
......@@ -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 = {
......
......@@ -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);
}
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册