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-05 16:50:40 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Update u2net_portrait_test.py

Browse Source
This commit is contained in:
adakoda 2020-12-24 11:04:43 +09:00 committed by GitHub
parent 21bc4de1f3
commit f893106780
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_test.py
View File

@ -84,10 +84,7 @@ def main():
print("...load U2NET---173.6 MB")
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()