Update README.md

README.md

@@ -43,11 +43,7 @@ tokenized = tokenizer(text, return_tensors='pt')
 output = model(**tokenized)
 
 # result
-logits = output['logits']
-logits = logits.detach().cpu().numpy()
-
-# remove start & end tag
-pred = [p for p in np.argmax(logits, axis=2)][0][1:-1]
+pred = np.argmax(output[0].cpu().detach().numpy(), axis=2)[0][1:-1]
 
 # Check pred
 for txt, pred in zip(tokenizer.tokenize(text), pred):