fix for 0.2

unos21 · Aug 6, 2017 · 9053040 · 9053040
1 parent 9012fae
commit 9053040
Showing 1 changed file with 1 addition and 6 deletions.
diff --git a/imagenet/main.py b/imagenet/main.py
@@ -16,11 +16,6 @@
 import torchvision.models as models
 
 
-model_names = sorted(name for name in models.__dict__
-    if name.islower() and not name.startswith("__")
-    and callable(models.__dict__[name]))
-
-
 parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
 parser.add_argument('data', metavar='DIR',
                     help='path to dataset')
@@ -308,7 +303,7 @@ def accuracy(output, target, topk=(1,)):
 
     res = []
     for k in topk:
-        correct_k = correct[:k].view(-1).float().sum(0)
+        correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
         res.append(correct_k.mul_(100.0 / batch_size))
     return res