From 1b247c5419c27cac6fd2595e4a4287e0888b4503 Mon Sep 17 00:00:00 2001 From: szerintedmi Date: Mon, 11 May 2020 12:01:26 +0100 Subject: [PATCH] load model to cpu when no cuda --- u2net_test.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/u2net_test.py b/u2net_test.py index 8b501c6..504ca55 100644 --- a/u2net_test.py +++ b/u2net_test.py @@ -83,9 +83,12 @@ def main(): elif(model_name=='u2netp'): print("...load U2NEP---4.7 MB") net = U2NETP(3,1) - net.load_state_dict(torch.load(model_dir)) + if torch.cuda.is_available(): + 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.eval() # --------- 4. inference for each image ---------