url valid

roatienza · May 3, 2021 · 035deda · 035deda
1 parent 91b9b1d
commit 035deda
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/requirements.txt b/requirements.txt
@@ -5,4 +5,5 @@ timm>=0.4.5
 lmdb>=1.0.0
 pillow>=8.2.0
 opencv-python>=4.5.1.48
-nltk>=3.4.5
+nltk>=3.4.5
+natsort>=7.1.0
diff --git a/test.py b/test.py
@@ -229,7 +229,11 @@ def test(opt):
 
     # load model
     print('loading pretrained model from %s' % opt.saved_model)
-    model.load_state_dict(torch.load(opt.saved_model, map_location=device))
+
+    if validators.url(opt.saved_model):
+        model.load_state_dict(torch.hub.load_state_dict_from_url(checkpoint, progress=True, map_location=device))
+    else:
+        model.load_state_dict(torch.load(opt.saved_model, map_location=device))
     opt.exp_name = '_'.join(opt.saved_model.split('/')[1:])
     # print(model)