diff options
| author | fban <fban@google.com> | 2022-08-10 22:01:32 +0000 |
|---|---|---|
| committer | Frank Ban <fban@google.com> | 2022-09-09 19:48:57 +0000 |
| commit | abaa5ec61b7e9ee979f9c2f1f596d4fa85bbabeb (patch) | |
| tree | 509068f9e113fd9a610a799163428e301cfa5dca | |
| parent | 589eb63fba6c01c7c7cf3cbf601f60d7b87334df (diff) | |
| download | tflite-support-abaa5ec61b7e9ee979f9c2f1f596d4fa85bbabeb.tar.gz | |
Adds a sanity check.
Missing this check was causing an ASAN error in G3. Adding it in AOSP as
well to be principled.
Bug: 241507692
Test: atest TfliteSupportClassifierTests
Change-Id: Ib0cf3893a08bace5584ede745324dd7bf50644ab
| -rw-r--r-- | tensorflow_lite_support/cc/task/text/nlclassifier/bert_nl_classifier.cc | 8 |
1 files changed, 8 insertions, 0 deletions
diff --git a/tensorflow_lite_support/cc/task/text/nlclassifier/bert_nl_classifier.cc b/tensorflow_lite_support/cc/task/text/nlclassifier/bert_nl_classifier.cc index d403755b..88a8a0af 100644 --- a/tensorflow_lite_support/cc/task/text/nlclassifier/bert_nl_classifier.cc +++ b/tensorflow_lite_support/cc/task/text/nlclassifier/bert_nl_classifier.cc @@ -246,6 +246,14 @@ absl::Status BertNLClassifier::InitializeFromMetadata() { TfLiteSupportStatus::kInvalidInputTensorSizeError); } + // If some tensor does not have a size 2 dims_signature, then we + // assume the input is not dynamic. + if (ids_tensor.dims_signature->size != 2 || + mask_tensor.dims_signature->size != 2 || + segment_ids_tensor.dims_signature->size != 2) { + return absl::OkStatus(); + } + if (ids_tensor.dims_signature->data[1] == -1 && mask_tensor.dims_signature->data[1] == -1 && segment_ids_tensor.dims_signature->data[1] == -1) { |