commit

a0405u · Dec 29, 2020 · 31d2902 · 31d2902
1 parent 0bc2346
commit 31d2902
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/lab3/lab3.py b/lab3/lab3.py
@@ -93,7 +93,8 @@ def predict(image):
     #output = model_trt(input)
     output = model(input)
 
-    probability = torch.nn.Softmax(output)
+    sm = torch.nn.Softmax()
+    probability = sm(output)
 
     print("image processed in {}s".format(round(time.time() - timest, 3)))