rjiang12 commited on
Commit
0647c43
·
1 Parent(s): 3eeacdb

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -64,7 +64,7 @@ def generate_answer_vilt(processor, model, image, question):
64
 
65
  with torch.no_grad():
66
  outputs = model(**encoding)
67
-
68
  return model.config.id2label[predicted_class_idx]
69
 
70
 
 
64
 
65
  with torch.no_grad():
66
  outputs = model(**encoding)
67
+ predicted_class_idx = outputs.logits.argmax(-1).item()
68
  return model.config.id2label[predicted_class_idx]
69
 
70