Update

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
2 years ago · 482a1caeff
1 changed files with 1 additions and 1 deletions
--- 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)