From 482a1caeffa2ed0b49ed79c8e02a08f894275d9d Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Tue, 16 Aug 2022 11:57:13 +0800 Subject: [PATCH] Update Signed-off-by: Jael Gu --- nn_fingerprint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nn_fingerprint.py b/nn_fingerprint.py index 2bb88c2..8ff3fcd 100644 --- a/nn_fingerprint.py +++ b/nn_fingerprint.py @@ -65,7 +65,7 @@ class NNFingerprint(NNOperator): log.warning('Using onnx.') self.model = onnxruntime.InferenceSession( model_path, - providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] + providers=['CPUExecutionProvider'] if self.device == 'cpu' else ['CUDAExecutionProvider'] ) else: state_dict = torch.load(model_path, map_location=self.device)