Hugging Face's logo

Spaces:
user-agent
/
background-remover
Running on Zero

App Files Community
user-agent commited on
Commit
8239fe8
·
verified ·
1 Parent(s): 98c6c6d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -13,7 +13,8 @@ torch.set_float32_matmul_precision("high")
13
  birefnet = AutoModelForImageSegmentation.from_pretrained(
14
  "ZhengPeng7/BiRefNet", trust_remote_code=True
15
  )
16
- birefnet.to("cuda")
 
17
 
18
  # Image transformation
19
  transform_image = transforms.Compose([
 
13
  birefnet = AutoModelForImageSegmentation.from_pretrained(
14
  "ZhengPeng7/BiRefNet", trust_remote_code=True
15
  )
16
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
17
+ birefnet.to(device)
18
 
19
  # Image transformation
20
  transform_image = transforms.Compose([