Fix passing sequence length and batch size to build the model.

PiperOrigin-RevId: 423714416

Fix passing sequence length and batch size to build the model.
PiperOrigin-RevId: 423714416
45848c44 · Renjie Liu · A. Unique TensorFlower · ab8b8012 · 45848c44
Commit 45848c44 authored Jan 23, 2022 by Renjie Liu Committed by A. Unique TensorFlower Jan 23, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

official/projects/edgetpu/nlp/serving/export_tflite_squad.py official/projects/edgetpu/nlp/serving/export_tflite_squad.py +2 -1

No files found.
--- a/official/projects/edgetpu/nlp/serving/export_tflite_squad.py
+++ b/official/projects/edgetpu/nlp/serving/export_tflite_squad.py
@@ -135,7 +135,8 @@ def main(argv: Sequence[str]) -> None:
    checkpoint = tf.train.Checkpoint(**checkpoint_dict)
    checkpoint.restore(FLAGS.model_checkpoint).assert_existing_objects_matched()

-  model_for_serving = build_model_for_serving(model)
+  model_for_serving = build_model_for_serving(model, FLAGS.sequence_length,
+                                              FLAGS.batch_size)
  model_for_serving.summary()

  # TODO(b/194449109): Need to save the model to file and then convert tflite