topk bug fix for binary classification

4 years ago · 87e2033b1d
parent 3cdaf5ed56
commit 87e2033b1d
1 changed files with 1 additions and 1 deletions
--- a/train.py
+++ b/train.py
@ -761,7 +761,7 @@ def validate(model, loader, loss_fn, args, amp_autocast=suppress, log_suffix='')
                target = target[0:target.size(0):reduce_factor]

            loss = loss_fn(output, target)
-            acc1, acc5 = accuracy(output, target, topk=(1, 5))
+            acc1, acc5 = accuracy(output, target, topk=(1, min(5, model.num_classes)))

            if args.distributed:
                reduced_loss = reduce_tensor(loss.data, args.world_size)