Remove module names from arch list in Imagenet example (pytorch#62)

RuiWang1998 · Jan 28, 2017 · 2cfe5e1 · 2cfe5e1
1 parent 9f125c1
commit 2cfe5e1
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/imagenet/main.py b/imagenet/main.py
@@ -15,7 +15,8 @@
 
 
 model_names = sorted(name for name in models.__dict__
-    if name.islower() and not name.startswith("__"))
+    if name.islower() and not name.startswith("__")
+    and callable(models.__dict__[name]))
 
 
 parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')