working changes

ManifoldRG · May 24, 2021 · 19ac706 · 19ac706
1 parent c5a8fed
commit 19ac706
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/models/efficientnet.py b/models/efficientnet.py
@@ -2,7 +2,7 @@
 import timm
 
 
-def load_efficientnet(model_name, num_classes=1000, input_channels=3, pretrained=False, ):
+def load_efficientnet(model_name="tf_efficientnet_b4", num_classes=1000, input_channels=3, pretrained=False):
 
  model = timm.create_model(
  model_name,

diff --git a/utils/helpers.py b/utils/helpers.py
@@ -103,7 +103,7 @@ def load_model(model_name, training_type, configs):
  model = vgg19_bn(in_channels=configs.input_channels,
  num_classes=configs.num_classes)
 
- elif model_name.startswith("EfficientNet"):
+ elif "efficientnet" in model_name:
 
  if training_type == 'pretrained':
  print(f"Loading pretrained {model_name}")