Commit cfdeaae6 authored by Mao, Bojia's avatar Mao, Bojia
Browse files

recorrect onnx to device

parent 658ccbca
......@@ -105,8 +105,10 @@ class ONNXInterface:
Returns :
label of model prediction and the corresponding confidence
'''
scoreFlag = False
if image_tensor != None:
scoreFlag = True
logits = self.onnx_model_forward(image_tensor,text_input)
else:
p = transforms.Compose([transforms.Scale((224,224))])
......@@ -124,7 +126,7 @@ class ONNXInterface:
scores = nn.functional.softmax(torch.tensor(logits), dim=1)
if image_tensor != None:
if scoreFlags:
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