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-17 12:35:52 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

map location to work on torch < 4.0.1

Browse Source
This commit is contained in:
szerintedmi 2020-05-11 18:10:21 +01:00
parent 1b247c5419
commit 4031385d80
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
u2net_test.py
View File

@ -88,7 +88,7 @@ def main():
net.load_state_dict(torch.load(model_dir)) net.load_state_dict(torch.load(model_dir))
net.cuda() net.cuda()
else: else:
net.load_state_dict(torch.load(model_dir, map_location=torch.device('cpu'))) net.load_state_dict(torch.load(model_dir, map_location='cpu'))
net.eval() net.eval()
# --------- 4. inference for each image --------- # --------- 4. inference for each image ---------