check if output is list

ln23415 · Feb 1, 2019 · 800df43 · 800df43
1 parent 564f5a2
commit 800df43
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/torchreid/utils/torchtools.py b/torchreid/utils/torchtools.py
@@ -84,6 +84,9 @@ def accuracy(output, target, topk=(1,)):
         maxk = max(topk)
         batch_size = target.size(0)
 
+        if isinstance(output, (tuple, list)):
+            output = output[0]
+
         _, pred = output.topk(maxk, 1, True, True)
         pred = pred.t()
         correct = pred.eq(target.view(1, -1).expand_as(pred))