diff --git a/test/backend/common.py b/test/backend/common.py index a9b077eae7..7e56f4ecb5 100644 --- a/test/backend/common.py +++ b/test/backend/common.py @@ -16,7 +16,15 @@ import sys import os import ctypes -import onnx + +try: + import onnx +except ImportError: + print( + "Failed to include the `onnx` package; please install to successfully use this script." + ) + exit(1) + import subprocess import variables from variables import * diff --git a/test/backend/inference_backend.py b/test/backend/inference_backend.py index b13f25f177..5bf36404cd 100644 --- a/test/backend/inference_backend.py +++ b/test/backend/inference_backend.py @@ -18,8 +18,16 @@ import base64 import numpy as np import re -import onnx -import onnx.parser + +try: + import onnx + import onnx.parser +except ImportError: + print( + "Failed to include the `onnx` package; please install to successfully use this script." + ) + exit(1) + import subprocess from onnx.backend.base import Device, DeviceType, Backend from onnx.backend.test import BackendTest diff --git a/test/backend/onnxmlir_node_tests.py b/test/backend/onnxmlir_node_tests.py index 9a63d34edc..864d3f767c 100644 --- a/test/backend/onnxmlir_node_tests.py +++ b/test/backend/onnxmlir_node_tests.py @@ -8,8 +8,16 @@ ################################################################################ import numpy as np -import onnx -import onnx.parser + +try: + import onnx + import onnx.parser +except ImportError: + print( + "Failed to include the `onnx` package; please install to successfully use this script." + ) + exit(1) + from collections import namedtuple OnnxMlirNodeTestCase = namedtuple(