commit

a0405u · Dec 29, 2020 · f0e4108 · f0e4108
1 parent 1837556
commit f0e4108
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/lab3/lab3.py b/lab3/lab3.py
@@ -94,7 +94,7 @@ def predict(image):
     output = model(input)
 
     sm = torch.nn.Softmax()
-    probability = sm(output)
+    probability = sm(output, dim = 1)
 
     print("image processed in {}s".format(round(time.time() - timest, 3)))