app

app.py

@@ -36,7 +36,7 @@ def mlm(image, text):
             with torch.no_grad():
                 predictions = model(tokens_tensor, segments_tensors)
             
-            _, predicted_index = torch.topk(predictions[0][0][masked_index], topk)#.item()
+            _, predicted_index = torch.topk(predictions[0][0][masked_index], 2)#.item()
             predicted_index = predicted_index.detach().cpu().numpy()
             #print(predicted_index)
             for idx in predicted_index: