chore: update encoding len

Anush008 · Jan 31, 2024 · df7e636 · df7e636
1 parent 770a12e
commit df7e636
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/fastembed.go b/fastembed.go
@@ -148,7 +148,7 @@ func (f *FlagEmbedding) onnxEmbed(input []string) ([]([]float32), error) {
 		inputTypeIdsFlat = append(inputTypeIdsFlat, inputTypeIds...)
 	}
 
-	inputShape := ort.NewShape(int64(len(inputs)), int64(f.maxLength))
+	inputShape := ort.NewShape(int64(len(inputs)), int64(encodings[0].Len()))
 
 	inputTensorID, err := ort.NewTensor(inputShape, inputIdsFlat)
 	if err != nil {
@@ -175,7 +175,7 @@ func (f *FlagEmbedding) onnxEmbed(input []string) ([]([]float32), error) {
 		return nil, err
 	}
 
-	outputShape := ort.NewShape(int64(len(inputs)), int64(f.maxLength), int64(modelInfo.Dim))
+	outputShape := ort.NewShape(int64(len(inputs)), int64(int64(encodings[0].Len())), int64(modelInfo.Dim))
 	outputTensor, err := ort.NewEmptyTensor[float32](outputShape)
 	if err != nil {
 		return nil, err