Update generic_ner.py

generic_ner.py

@@ -234,8 +234,8 @@ class MultitaskTokenClassificationPipeline(Pipeline):
         predictions = {}
         confidence_scores = {}
         for task, logits in tokens_result.logits.items():
-            predictions[task] = torch.argmax(logits, dim=-1).tolist()
-            confidence_scores[task] = F.softmax(logits, dim=-1).tolist()
+            predictions[task] = torch.argmax(logits, dim=-1).tolist()[0]
+            confidence_scores[task] = F.softmax(logits, dim=-1).tolist()[0]
 
         decoded_predictions = {}
         for task, preds in predictions.items():