少量修改。

LIU42 · Sep 5, 2024 · 1dd46d2 · 1dd46d2
1 parent 652ab83
commit 1dd46d2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/classify/predict.py b/classify/predict.py
@@ -6,7 +6,7 @@
 
 class ShapeClassifier:
     def __init__(self, configs):
-        if configs['device'] == 'GPU':
+        if configs['device'] == 'CUDA':
             providers = ['CUDAExecutionProvider']
         else:
             providers = ['CPUExecutionProvider']