fix import

Author: xadupre

_doc/examples/plot_failing_model_extract.py

@@ -27,7 +27,7 @@
 import onnxruntime
 from onnx_diagnostic import doc
 from onnx_diagnostic.helpers.onnx_helper import from_array_extended
-from onnx_diagnostic.ort_session import investigate_onnxruntime_issue
+from onnx_diagnostic.helpers.ort_session import investigate_onnxruntime_issue
 
 TFLOAT = onnx.TensorProto.FLOAT
 

_doc/examples/plot_failing_onnxruntime_evaluator.py

@@ -11,9 +11,9 @@
 It splits the model into node and runs them independently until it succeeds
 or fails. This class converts every node into model based on the types
 discovered during the execution. It relies on :class:`InferenceSessionForTorch
-<onnx_diagnostic.ort_session.InferenceSessionForTorch>` or
+<onnx_diagnostic.helpers.ort_session.InferenceSessionForTorch>` or
 :class:`InferenceSessionForNumpy
-<onnx_diagnostic.ort_session.InferenceSessionForNumpy>`
+<onnx_diagnostic.helpers.ort_session.InferenceSessionForNumpy>`
 for the execution. This example uses torch tensor and
 bfloat16.
 

_unittests/ut_torch_onnx/test_sbs.py

@@ -1,5 +1,5 @@
 import unittest
-from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout
+from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, ignore_warnings
 from onnx_diagnostic.reference import ExtendedReferenceEvaluator
 from onnx_diagnostic.torch_onnx.sbs import run_aligned
 
@@ -13,6 +13,7 @@ class TestSideBySide(ExtTestCase):
 
     @hide_stdout()
     @unittest.skipIf(to_onnx is None, "to_onnx not installed")
+    @ignore_warnings((UserWarning,))
     def test_ep_onnx_sync_exp(self):
         import torch
 
@@ -44,8 +45,9 @@ def forward(self, x):
         self.assertEqual(len(results), 4)
 
     @hide_stdout()
+    @ignore_warnings((DeprecationWarning, FutureWarning, UserWarning))
     @unittest.skipIf(to_onnx is None, "to_onnx not installed")
-    def test_ep_onnx_sync(self):
+    def test_ep_onnx_sync_a(self):
         import torch
 
         class Model(torch.nn.Module):

onnx_diagnostic/ext_test_case.py

@@ -1044,7 +1044,7 @@ def assert_onnx_disc(
         :param atol: absolute tolerance
         :param rtol: relative tolerance
         :param kwargs: arguments sent to
-            :class:`onnx_diagnostic.ort_session.InferenceSessionForTorch`
+            :class:`onnx_diagnostic.helpers.ort_session.InferenceSessionForTorch`
         """
         import torch.utils._pytree as py_pytree
         from .helper import string_type, string_diff, max_diff

onnx_diagnostic/helpers/ort_session.py

@@ -522,7 +522,7 @@ def investigate_onnxruntime_issue(
         import numpy as np
         import onnx
         import onnx.helper as oh
-        from onnx_diagnostic.ort_session import investigate_onnxruntime_issue
+        from onnx_diagnostic.helpers.ort_session import investigate_onnxruntime_issue
 
         TFLOAT = onnx.TensorProto.FLOAT
         model = oh.make_model(