From 95a92cb092bbfb5e63908d6cefd7ae3d12a2f729 Mon Sep 17 00:00:00 2001 From: fangrenfu Date: Fri, 13 Mar 2020 00:37:13 +0800 Subject: [PATCH] fix 0-dim tensor error Use tensor.item() to convert a 0-dim tensor to a Python number --- examples/imagenet_logits.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/imagenet_logits.py b/examples/imagenet_logits.py index ef2cbf1..fe05cb9 100644 --- a/examples/imagenet_logits.py +++ b/examples/imagenet_logits.py @@ -62,7 +62,7 @@ def main(): # Make predictions output = model(input) # size(1, 1000) max, argmax = output.data.squeeze().max(0) - class_id = argmax[0] + class_id = argmax.item() class_key = class_id_to_key[class_id] classname = key_to_classname[class_key]