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_test.py

Browse Source
This commit is contained in:
adakoda 2020-12-24 11:13:03 +09:00 committed by GitHub
parent 6b2d7be4fd
commit 4ea81c2f21
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_test.py
View File

@ -84,10 +84,7 @@ def main():
elif(model_name=='u2netp'):
print("...load U2NEP---4.7 MB")
net = U2NETP(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()