| 10 | |
| 11 | namespace { |
| 12 | cg::OperatorNodeBase* apply_on_var_node(const OpDef& def, const VarNodeArray& inputs) { |
| 13 | auto&& nms_keep = def.cast_final_safe<NMSKeep>(); |
| 14 | |
| 15 | NMSKeepOpr::Param param; |
| 16 | param.iou_thresh = nms_keep.iou_thresh; |
| 17 | param.max_output = nms_keep.max_output; |
| 18 | |
| 19 | OperatorNodeConfig config{nms_keep.make_name()}; |
| 20 | |
| 21 | return NMSKeepOpr::make(inputs[0], param, config).node()->owner_opr(); |
| 22 | } |
| 23 | |
| 24 | OP_TRAIT_REG(NMSKeep, NMSKeep, NMSKeepOpr) |
| 25 | .apply_on_var_node(apply_on_var_node) |