From 3e0885f88da328aefb7abb2fa350f8860a4bd52d Mon Sep 17 00:00:00 2001
From: andong <donganchn@163.com>
Date: Thu, 19 Dec 2019 18:36:22 +0800
Subject: [PATCH] correctly print model_names
---
dirtorch/nets/__init__.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/dirtorch/nets/__init__.py b/dirtorch/nets/__init__.py
index 1df2bda..29d844d 100644
--- 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)
--
GitLab