user-agent commited on
Commit
1e22887
·
verified ·
1 Parent(s): 8239fe8

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -25,7 +25,7 @@ transform_image = transforms.Compose([
25
 
26
  def process(image: Image.Image) -> Image.Image:
27
  image_size = image.size
28
- input_tensor = transform_image(image).unsqueeze(0).to("cuda")
29
 
30
  with torch.no_grad():
31
  preds = birefnet(input_tensor)[-1].sigmoid().cpu()
 
25
 
26
  def process(image: Image.Image) -> Image.Image:
27
  image_size = image.size
28
+ input_tensor = transform_image(image).unsqueeze(0).to(device)
29
 
30
  with torch.no_grad():
31
  preds = birefnet(input_tensor)[-1].sigmoid().cpu()