Update app.py

app.py

@@ -67,7 +67,7 @@ def generate_answer_vilt(processor, model, image, question):
     lsm = torch.nn.LogSoftmax(dim=1)
     print(lsm(outputs.logits))
     predicted_class_idx = outputs.logits.argmax(-1).item()
-    logitsList = outputs.logits.tolist()
+    logitsList = outputs.logits.flatten().tolist()
     print(logitsList)
     maybeProbsList = [math.exp(i) for i in logitsList]
     return model.config.id2label[predicted_class_idx]