From 4031385d809406d54779d6d934d99fd2c5d988b8 Mon Sep 17 00:00:00 2001 From: szerintedmi Date: Mon, 11 May 2020 18:10:21 +0100 Subject: [PATCH] map location to work on torch < 4.0.1 --- u2net_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/u2net_test.py b/u2net_test.py index 504ca55..8c33886 100644 --- a/u2net_test.py +++ b/u2net_test.py @@ -88,7 +88,7 @@ def main(): net.load_state_dict(torch.load(model_dir)) net.cuda() 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() # --------- 4. inference for each image ---------