Tweak accuracy topk safety. Fix #807

pull/828/head
Ross Wightman 3 years ago
parent 35c9740826
commit d667351eac

@ -2,7 +2,6 @@
Hacked together by / Copyright 2020 Ross Wightman
"""
import torch
class AverageMeter:
@ -30,7 +29,4 @@ def accuracy(output, target, topk=(1,)):
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
correct = pred.eq(target.reshape(1, -1).expand_as(pred))
return [
correct[:k].reshape(-1).float().sum(0) * 100. / batch_size
if k <= maxk else torch.tensor(100.) for k in topk
]
return [correct[:min(k, maxk)].reshape(-1).float().sum(0) * 100. / batch_size for k in topk]

Loading…
Cancel
Save