adds missing method

Author: xadupre

_unittests/ut_reference/test_ort_evaluator.py

@@ -20,6 +20,7 @@
     onnx_dtype_to_np_dtype,
 )
 from onnx_diagnostic.reference import ExtendedReferenceEvaluator, OnnxruntimeEvaluator
+from onnx_diagnostic.helpers.ort_session import _InferenceSession
 
 TFLOAT = TensorProto.FLOAT
 
@@ -78,6 +79,18 @@ def test_ort_eval(self):
         self.assertEqualArray(expected, got, atol=1e-4)
         self.assertIn("Reshape(xm, shape3) -> Z", out)
 
+    @ignore_warnings(DeprecationWarning)
+    def test__inference(self):
+        model = self._get_model()
+
+        feeds = {"X": self._range(32, 128), "Y": self._range(3, 5, 128, 64)}
+        ref = ExtendedReferenceEvaluator(model)
+        expected = ref.run(None, feeds)[0]
+
+        ort_eval = _InferenceSession(model)
+        got = ort_eval.run(None, feeds)[0]
+        self.assertEqualArray(expected, got, atol=1e-4)
+
     @ignore_warnings(DeprecationWarning)
     @requires_cuda()
     @hide_stdout()

onnx_diagnostic/helpers/ort_session.py

@@ -155,6 +155,19 @@ def __init__(
 
         self._torch_from_dlpack = _from_dlpack
 
+    def run(
+        self,
+        output_names: Optional[List[str]],
+        feeds: Union[Dict[str, np.ndarray], Dict[str, "OrtValue"]],  # noqa: F821
+    ) -> Union[List[np.ndarray], List["OrtValue"]]:  # noqa: F821
+        """Calls :meth:`onnxruntime.InferenceSession.run`."""
+        if any(isinstance(t, np.ndarray) for t in feeds.values()):
+            return self.sess.run(output_names, feeds)
+        ort_outputs = self.sess._sess.run_with_ort_values(
+            feeds, output_names or self.output_names, self.run_options
+        )
+        return ort_outputs
+
 
 class InferenceSessionForNumpy(_InferenceSession):
     """