unit test

microsoft · Apr 1, 2024 · f876652 · f876652
1 parent ea0a057
commit f876652
Showing 1 changed file with 20 additions and 0 deletions.
diff --git a/onnxruntime/test/python/onnxruntime_test_python.py b/onnxruntime/test/python/onnxruntime_test_python.py
@@ -354,6 +354,26 @@ def test_set_providers_with_options(self):
                 sess.set_providers(['TensorrtExecutionProvider'], [option])
             """
 
+            # test for user_compute_stream
+            option["user_compute_stream"] = "0"
+            sess.set_providers(["TensorrtExecutionProvider"], [option])
+            options = sess.get_provider_options()
+            self.assertEqual(options["TensorrtExecutionProvider"]["user_compute_stream"], "0")
+
+            try:
+                import torch
+
+                if torch.cuda.is_available():
+                    s = torch.cuda.Stream()
+                    option["user_compute_stream"] = str(s.cuda_stream)
+                    sess.set_providers(["TensorrtExecutionProvider"], [option])
+                    options = sess.get_provider_options()
+                    self.assertEqual(options["TensorrtExecutionProvider"]["user_compute_stream"], str(s.cuda_stream))
+                    self.assertEqual(options["TensorrtExecutionProvider"]["has_user_compute_stream"], "1")
+            except ImportError:
+                print("torch is not installed, skip testing setting user_compute_stream from torch cuda stream")
+
+
         if "CUDAExecutionProvider" in onnxrt.get_available_providers():
             cuda_success = 0