Hugging Face's logo

Spaces:
hylee
/
u2net_portrait
Build error

App Files Community
hylee commited on
Commit
c388d64
•
1 Parent(s): ca61862

init

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -75,7 +75,7 @@ model_dir = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'U-2-Net/sa
75
  print("...load U2NET---173.6 MB")
76
  net = U2NET(3,1)
77
 
78
- net.load_state_dict(torch.load(model_dir))
79
  # if torch.cuda.is_available():
80
  # net.cuda()
81
  net.eval()
 
75
  print("...load U2NET---173.6 MB")
76
  net = U2NET(3,1)
77
 
78
+ net.load_state_dict(torch.load(model_dir, map_location='cpu'))
79
  # if torch.cuda.is_available():
80
  # net.cuda()
81
  net.eval()