Merge pull request #69 from sunhailin-Leo/dev

fix missing variable for grpc request
sunhailin-Leo · May 27, 2024 · fbc830d · fbc830d
2 parents 411b960 + fdb012c
commit fbc830d
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/models/transformers/bert.go b/models/transformers/bert.go
@@ -233,18 +233,21 @@ func (m *BertModelService) generateGRPCRequest(
 					m.grpcSliceToLittleEndianByteSlice(
 						m.MaxSeqLength, feature.TypeIDs, inferInputTensor[j].Datatype)...,
 				)
+				inferInputTensor[j].Shape = []int64{int64(len(inferDataArr)), int64(m.MaxSeqLength)}
 			case ModelBertModelInputIdsKey:
 				inputIdsBytes = append(
 					inputIdsBytes,
 					m.grpcSliceToLittleEndianByteSlice(
 						m.MaxSeqLength, feature.TokenIDs, inferInputTensor[j].Datatype)...,
 				)
+				inferInputTensor[j].Shape = []int64{int64(len(inferDataArr)), int64(m.MaxSeqLength)}
 			case ModelBertModelInputMaskKey:
 				inputMaskBytes = append(
 					inputMaskBytes,
 					m.grpcSliceToLittleEndianByteSlice(
 						m.MaxSeqLength, feature.Mask, inferInputTensor[j].Datatype)...,
 				)
+				inferInputTensor[j].Shape = []int64{int64(len(inferDataArr)), int64(m.MaxSeqLength)}
 			}
 		}
 		batchModelInputObjs[i] = inputObject