[PHI] Add multiclass_nms3 output defs (#51355)

* add nms3 register output defs * remove nms from set * remove nms from set

[PHI] Add multiclass_nms3 output defs (#51355)
* add nms3 register output defs * remove nms from set * remove nms from set
06cb6553 · PuQing · GitHub · 289677e2 · 06cb6553
隐藏空白更改
内联并排

Showing with 2 addition and 0 deletion

paddle/phi/kernels/cpu/multiclass_nms3_kernel.cc paddle/phi/kernels/cpu/multiclass_nms3_kernel.cc +2 -0

未找到文件。
--- a/paddle/phi/kernels/cpu/multiclass_nms3_kernel.cc
+++ b/paddle/phi/kernels/cpu/multiclass_nms3_kernel.cc
@@ -628,4 +628,6 @@ void MultiClassNMSKernel(const Context& ctx,

 PD_REGISTER_KERNEL(
    multiclass_nms3, CPU, ALL_LAYOUT, phi::MultiClassNMSKernel, float, double) {
+  kernel->OutputAt(1).SetDataType(phi::DataType::INT32);
+  kernel->OutputAt(2).SetDataType(phi::DataType::INT32);
 }