diff options
-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) { |