commit

a0405u · Dec 29, 2020 · 6f43ecd · 6f43ecd
1 parent 1a095ef
commit 6f43ecd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/lab3/lab3.py b/lab3/lab3.py
@@ -97,7 +97,7 @@ def predict(image):
 
     print("image processed in {}s".format(round(time.time() - timest, 3)))
 
-    return output.data.cpu().numpy().argmax(), probability
+    return output.data.cpu().numpy().argmax(), probability[output.data.cpu().numpy().argmax()]
 
 
 # Process image