From 16f0e087a375bee6d2f56c090255bc8a454eebeb Mon Sep 17 00:00:00 2001
From: Chi Lo <Chi.Lo@microsoft.com>
Date: Thu, 30 Jan 2025 00:58:51 +0000
Subject: [PATCH] sort output to pass NMS unit test

---
 .../non_max_suppression_test.cc               | 24 +++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

diff --git a/onnxruntime/test/providers/cpu/object_detection/non_max_suppression_test.cc b/onnxruntime/test/providers/cpu/object_detection/non_max_suppression_test.cc
index 4b8475e3f6..f1180735b7 100644
--- a/onnxruntime/test/providers/cpu/object_detection/non_max_suppression_test.cc
+++ b/onnxruntime/test/providers/cpu/object_detection/non_max_suppression_test.cc
@@ -63,13 +63,23 @@ TEST(NonMaxSuppressionOpTest, TwoClasses) {
   test.AddInput<int64_t>("max_output_boxes_per_class", {}, {6L});
   test.AddInput<float>("iou_threshold", {}, {0.5f});
   test.AddInput<float>("score_threshold", {}, {0.0f});
+
+// The selected_indices in ORT is sorted by class, whereas in TRT the selected_indices is sorted by score,
+// So it needs to sort the output to pass the output check when there is more than one class using TRT EP.
+#ifdef USE_TENSORRT
+  bool sort_output = true;
+#else
+  bool sort_output = false;  // default
+#endif
+
   test.AddOutput<int64_t>("selected_indices", {6, 3},
                           {0L, 0L, 3L,
                            0L, 0L, 0L,
                            0L, 0L, 5L,
                            0L, 1L, 3L,
                            0L, 1L, 0L,
-                           0L, 1L, 5L});
+                           0L, 1L, 5L},
+                          sort_output);
   test.Run();
 }
 
@@ -125,6 +135,15 @@ TEST(NonMaxSuppressionOpTest, TwoBatches_TwoClasses) {
                         0.1f, 0.2f, 0.6f, 0.3f, 0.9f});
   test.AddInput<int64_t>("max_output_boxes_per_class", {}, {2L});
   test.AddInput<float>("iou_threshold", {}, {0.8f});
+
+// The selected_indices in ORT is sorted by class, whereas in TRT the selected_indices is sorted by score,
+// So it needs to sort the output to pass the output check when there is more than one class using TRT EP.
+#ifdef USE_TENSORRT
+  bool sort_output = true;
+#else
+  bool sort_output = false;  // default
+#endif
+
   test.AddOutput<int64_t>("selected_indices", {8, 3},
                           {0L, 0L, 4L,
                            0L, 0L, 2L,
@@ -134,7 +153,8 @@ TEST(NonMaxSuppressionOpTest, TwoBatches_TwoClasses) {
                            1L, 0L, 4L,
                            1L, 0L, 1L,
                            1L, 1L, 4L,
-                           1L, 1L, 1L});
+                           1L, 1L, 1L},
+                          sort_output);
   test.Run();
 }