correctly print model_names

dongan-beta · Dec 19, 2019 · 3e0885f · 3e0885f
1 parent 5bcb9fc
commit 3e0885f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dirtorch/nets/__init__.py b/dirtorch/nets/__init__.py
@@ -34,7 +34,7 @@ def create_model(arch, pretrained='', delete_fc=False, *args, **kwargs):
         optional arguments
     '''
     # creating model
-    if arch not in globals():
+    if arch not in model_names:
         raise NameError("unknown model architecture '%s'\nSelect one in %s" % (
                          arch, ','.join(model_names)))
     model = globals()[arch](*args, **kwargs)