Support downloading SwinIR pretrained models

cszn · Sep 9, 2021 · 7b4b55f · 7b4b55f
1 parent 6c38f59
commit 7b4b55f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/main_download_pretrained_models.py b/main_download_pretrained_models.py
@@ -111,7 +111,7 @@ def download_pretrained_model(model_dir='model_zoo', model_name='dncnn3.pth'):
                     else:
                         download_pretrained_model(args.model_dir, model_name)
             elif method_model in model_zoo:  # model, do not need for loop
-                if 'SwinIR' in model_name:
+                if 'SwinIR' in method_model:
                     download_pretrained_model(os.path.join(args.model_dir, 'swinir'), method_model)
                 else:
                     download_pretrained_model(args.model_dir, method_model)
@@ -125,4 +125,4 @@ def download_pretrained_model(model_dir='model_zoo', model_name='dncnn3.pth'):
 
 
 
-
+