diff --git a/timm/utils/metrics.py b/timm/utils/metrics.py index 4f5d95a1..9fdbe13e 100644 --- a/timm/utils/metrics.py +++ b/timm/utils/metrics.py @@ -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]