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

recorrect onnx

parent a05f95dd
......@@ -71,7 +71,7 @@ class ONNXInterface(torch.nn.Module):
if count == 3:
print("Assume it is bert model with only text input")
tokens = Tokenizer(text_input, return_tensors="pt")
ort_inputs = {k: v.cpu().detach().numpy() for k, v in token.items()}
ort_inputs = {k: v.cpu().detach().numpy() for k, v in tokens.items()}
elif count == 2:
print("Assume image and one text")
......
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