From f5a7a53dbd303c7d02706f289871d95514cf5e8f Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Sat, 13 Jan 2024 23:31:43 +0900 Subject: [PATCH] olive 0.4.x support --- modules/onnx_pipelines.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/modules/onnx_pipelines.py b/modules/onnx_pipelines.py index 86abb1439..2a0a25722 100644 --- a/modules/onnx_pipelines.py +++ b/modules/onnx_pipelines.py @@ -146,7 +146,10 @@ class OnnxRawPipeline(OnnxPipelineBase): try: from olive.workflows import run - from olive.model import ONNXModel + try: + from olive.model import ONNXModel + except ImportError: + from olive.model import ONNXModelHandler as ONNXModel shutil.rmtree("cache", ignore_errors=True) shutil.rmtree("footprints", ignore_errors=True) @@ -252,7 +255,10 @@ class OnnxRawPipeline(OnnxPipelineBase): try: from olive.workflows import run - from olive.model import ONNXModel + try: + from olive.model import ONNXModel + except ImportError: + from olive.model import ONNXModelHandler as ONNXModel shutil.rmtree("cache", ignore_errors=True) shutil.rmtree("footprints", ignore_errors=True)