Merge pull request #64 from MichelML/patch-1

Improved error messages

Author: lukemelas

efficientnet_pytorch/model.py

@@ -207,6 +207,6 @@ def _check_model_name_is_valid(cls, model_name, also_need_pretrained_weights=Fal
         """ Validates model name. None that pretrained weights are only available for
         the first four models (efficientnet-b{i} for i in 0,1,2,3) at the moment. """
         num_models = 4 if also_need_pretrained_weights else 8
-        valid_models = ['efficientnet_b'+str(i) for i in range(num_models)]
-        if model_name.replace('-','_') not in valid_models:
+        valid_models = ['efficientnet-b'+str(i) for i in range(num_models)]
+        if model_name not in valid_models:
             raise ValueError('model_name should be one of: ' + ', '.join(valid_models))