From 12945145d36c32b4968d12c552095580d0527432 Mon Sep 17 00:00:00 2001 From: john Date: Tue, 5 Dec 2023 19:25:56 -0600 Subject: [PATCH] close https://github.com/nadermx/backgroundremover/issues/112 --- backgroundremover/bg.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/backgroundremover/bg.py b/backgroundremover/bg.py index ffa4273..388e8cc 100644 --- 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):