Merge pull request #137 from lukemelas/advprop

Fixed `url_map` name
lukemelas · Jan 24, 2020 · f8d1eb4 · f8d1eb4
2 parents 50b0de9 + 53dbbe5
commit f8d1eb4
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/efficientnet_pytorch/__init__.py b/efficientnet_pytorch/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "0.6.0"
+__version__ = "0.6.1"
 from .model import EfficientNet
 from .utils import (
     GlobalParams,

diff --git a/efficientnet_pytorch/utils.py b/efficientnet_pytorch/utils.py
@@ -295,7 +295,7 @@ def get_model_params(model_name, override_params):
     return blocks_args, global_params
 
 
-url_map_aa = {
+url_map = {
     'efficientnet-b0': 'https://publicmodels.blob.core.windows.net/container/aa/efficientnet-b0-355c32eb.pth',
     'efficientnet-b1': 'https://publicmodels.blob.core.windows.net/container/aa/efficientnet-b1-f1951068.pth',
     'efficientnet-b2': 'https://publicmodels.blob.core.windows.net/container/aa/efficientnet-b2-8bb594d6.pth',
@@ -323,8 +323,8 @@ def get_model_params(model_name, override_params):
 def load_pretrained_weights(model, model_name, load_fc=True, advprop=False):
     """ Loads pretrained weights, and downloads if loading for the first time. """
     # AutoAugment or Advprop (different preprocessing)
-    url_map = url_map_advprop if advprop else url_map_aa
-    state_dict = model_zoo.load_url(url_map[model_name])
+    url_map_ = url_map_advprop if advprop else url_map
+    state_dict = model_zoo.load_url(url_map_[model_name])
     if load_fc:
         model.load_state_dict(state_dict)
     else:

diff --git a/setup.py b/setup.py
@@ -18,7 +18,7 @@
 EMAIL = 'lmelaskyriazi@college.harvard.edu'
 AUTHOR = 'Luke'
 REQUIRES_PYTHON = '>=3.5.0'
-VERSION = '0.6.0'
+VERSION = '0.6.1'
 
 # What packages are required for this module to be executed?
 REQUIRED = [