diff --git a/timm/utils/metrics.py b/timm/utils/metrics.py index 8e0b1f99..9fdbe13e 100644 --- a/timm/utils/metrics.py +++ b/timm/utils/metrics.py @@ -24,9 +24,9 @@ class AverageMeter: def accuracy(output, target, topk=(1,)): """Computes the accuracy over the k top predictions for the specified values of k""" - maxk = max(topk) + maxk = min(max(topk), output.size()[1]) batch_size = target.size(0) _, 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 for k in topk] + return [correct[:min(k, maxk)].reshape(-1).float().sum(0) * 100. / batch_size for k in topk]