Fix NonMaxSuppression default values in ReferenceEvaluator (onnx#6354)

### Description Fixes a bug in ReferenceEvaluator. --------- Signed-off-by: xadupre <[email protected]> Co-authored-by: Andreas Fehlner <[email protected]>
georgen117 · Sep 9, 2024 · 976a142 · 976a142
1 parent acb2c0a
commit 976a142
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git a/onnx/reference/custom_element_types.py b/onnx/reference/custom_element_types.py
@@ -46,7 +46,7 @@ def convert_from_ml_dtypes(array: np.ndarray) -> np.ndarray:
     if not ml_dtypes:
         return array
     for dtype, _, ml_name in _supported_types:
-        if array.dtype == getattr(ml_dtypes, ml_name):
+        if array.dtype == getattr(ml_dtypes, ml_name, None):
             return array.view(dtype=dtype)
     return array
 

diff --git a/onnx/reference/ops/op_non_max_suppression.py b/onnx/reference/ops/op_non_max_suppression.py
@@ -160,7 +160,7 @@ def prepare_compute(  # type: ignore
             pc.max_output_boxes_per_class_ = max_output_boxes_per_class_tensor
         if iou_threshold_tensor.size != 0:
             pc.iou_threshold_ = iou_threshold_tensor
-        if score_threshold_tensor.size != 0:
+        if score_threshold_tensor is not None and score_threshold_tensor.size != 0:
             pc.score_threshold_ = score_threshold_tensor
 
         pc.boxes_size_ = boxes_tensor.size
@@ -177,13 +177,11 @@ def _run(  # type: ignore
         self,
         boxes,
         scores,
-        max_output_boxes_per_class,
-        iou_threshold,
-        score_threshold,
-        center_point_box,
+        max_output_boxes_per_class=None,
+        iou_threshold=None,
+        score_threshold=None,
+        center_point_box=None,
     ):
-        center_point_box = center_point_box or self.center_point_box  # type: ignore
-
         pc = PrepareContext()
         self.prepare_compute(
             pc,