AI/U-2-Net
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/xuebinqin/U-2-Net synced 2025-08-06 05:39:57 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Update u2net_portrait_demo.py

Browse Source
This commit is contained in:
adakoda 2020-12-24 11:03:29 +09:00 committed by GitHub
parent db4815f38d
commit 21bc4de1f3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
u2net_portrait_demo.py
View File

@ -151,10 +151,7 @@ def main():
# load u2net_portrait model
net = U2NET(3,1)
if torch.cuda.is_available():
net.load_state_dict(torch.load(model_dir))
else:
net.load_state_dict(torch.load(model_dir, map_location=torch.device('cpu')))
net.load_state_dict(torch.load(model_dir))
if torch.cuda.is_available():
net.cuda()
net.eval()