diff --git a/models/transformers/bert_w2ner.go b/models/transformers/bert_w2ner.go
index 53f0f95..bf99f5e 100644
--- a/models/transformers/bert_w2ner.go
+++ b/models/transformers/bert_w2ner.go
@@ -111,15 +111,11 @@ func (w *W2NerModelService) getBertInputFeature(batchInferData [][]string) []*W2
 		for j := 0; j < padTokenLen; j++ {
 			gridMask2d[j] = generateAllTrueSlice(padTokenLen)
 
-			if j >= len(inferTokens) {
+			if j >= len(inferTokens) || len(inferTokens[j]) == 0 {
 				pieces2word[j] = make([]bool, len(batchInputFeatures[i].TokenIDs))
 				continue
 			}
 
-			if len(inferTokens[j]) == 0 {
-				continue
-			}
-
 			idx := utils.GenerateRange[int](start, start+len(inferTokens[j]))
 			pieces2word[j] = make([]bool, len(batchInputFeatures[i].TokenIDs))
 			for k := 0; k < len(pieces2word[j]); k++ {