Update README.md

README.md

@@ -43,13 +43,9 @@ tokenized = tokenizer(text, return_tensors='pt')
 output = model(**tokenized)
 
 # result
-logits = output['logits']
-logits = logits.detach().cpu().numpy()
+pred = np.argmax(output[0].cpu().detach().numpy(), axis=2)[0][1:-1]
 
-# remove start & end tag
-pred = [p for p in np.argmax(logits, axis=2)][0][1:-1]
-
-# Check pred
+# check pred
 for txt, pred in zip(tokenizer.tokenize(text), pred):
     print("{}\t{}".format(id2tag[pred], txt))
     # B_MT mental