aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--tensorflow_lite_support/cc/task/text/nlclassifier/bert_nl_classifier.cc8
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) {