Commit 8e43d94f authored by Mao, Bojia's avatar Mao, Bojia
Browse files

recorrect onnx to device

parent cfdeaae6
......@@ -126,7 +126,7 @@ class ONNXInterface:
scores = nn.functional.softmax(torch.tensor(logits), dim=1)
if scoreFlags:
if scoreFlag == True:
return scores
confidence, label = torch.max(scores, dim=1)
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment