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

recorrect onnx to device

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