fix device selection, increase version (#756)

obss · Nov 27, 2022 · daeef7e · daeef7e
1 parent 7838001
commit daeef7e
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/sahi/__init__.py b/sahi/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "0.11.4"
+__version__ = "0.11.5"
 
 from sahi.annotation import BoundingBox, Category, Mask
 from sahi.auto_model import AutoDetectionModel

diff --git a/sahi/utils/torch.py b/sahi/utils/torch.py
@@ -64,7 +64,7 @@ def select_device(device: str):
     """
     if device == "cuda":
         device = "cuda:0"
-    device = str(device).strip().lower().replace("none", "")  # to string, 'cuda:0' to '0'
+    device = str(device).strip().lower().replace("cuda:", "").replace("none", "")  # to string, 'cuda:0' to '0'
     cpu = device == "cpu"
     mps = device == "mps"  # Apple Metal Performance Shaders (MPS)
     if cpu or mps: