From d4b82c7f9b4222a979029a6d75756cf43e895101 Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Tue, 13 Dec 2022 17:00:19 +0800 Subject: [PATCH] Update Signed-off-by: Jael Gu --- nn_fingerprint.py | 2 +- requirements.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/nn_fingerprint.py b/nn_fingerprint.py index 6eadb50..131ab36 100644 --- a/nn_fingerprint.py +++ b/nn_fingerprint.py @@ -22,7 +22,6 @@ from typing import List, Union import torch import torchaudio import numpy -import onnxruntime from towhee.operator.base import NNOperator from towhee import register @@ -67,6 +66,7 @@ class NNFingerprint(NNOperator): model_path = os.path.join(path, 'saved_model', 'nnfp_fma.pt') if model_path.endswith('.onnx'): log.warning('Using onnx.') + import onnxruntime self.model = onnxruntime.InferenceSession( model_path, providers=['CPUExecutionProvider'] if self.device == 'cpu' else ['CUDAExecutionProvider'] diff --git a/requirements.txt b/requirements.txt index 40f1535..d0d939d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ towhee.models>=0.8.0 torch torchaudio -onnxruntime +# onnxruntime