Update model_wrapper/bert_wrapper.py

model_wrapper/bert_wrapper.py

@@ -34,4 +34,4 @@ class BertWrapper(object):
         )  # output is logits for huggingfcae transformers
         predicted = torch.nn.functional.softmax(outputs.logits, dim=1)
         predicted_id = torch.argmax(predicted, dim=1).numpy()[0]
-        return self.id2label[predicted_id], predicted[predicted_id]
+        return self.id2label[predicted_id], predicted[0][predicted_id]