close #112

nadermx · Dec 6, 2023 · 1294514 · 1294514
1 parent ec7c796
commit 1294514
Showing 1 changed file with 10 additions and 5 deletions.
diff --git a/backgroundremover/bg.py b/backgroundremover/bg.py
@@ -16,11 +16,16 @@
 from . import utilities
 
 # closes https://github.com/nadermx/backgroundremover/issues/18
-if torch.cuda.is_available():
-    DEVICE = torch.device('cuda:0')
-elif torch.backends.mps.is_available():
-    DEVICE = torch.device('mps')
-else:
+# closes https://github.com/nadermx/backgroundremover/issues/112
+try:
+    if torch.cuda.is_available():
+        DEVICE = torch.device('cuda:0')
+    elif torch.backends.mps.is_available():
+        DEVICE = torch.device('mps')
+    else:
+        DEVICE = torch.device('cpu')
+except Exception as e:
+    print(f"An error occurred: {e}")
     DEVICE = torch.device('cpu')
 
 class Net(torch.nn.Module):