Skip to content

GitLab

Commit 25b9330e authored by Renjie Liu's avatar Renjie Liu Committed by A. Unique TensorFlower
Browse files

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

PiperOrigin-RevId: 423714416
parent a2fc0302
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
official/projects/edgetpu/nlp/serving/export_tflite_squad.py
View file @ 25b9330e
...@@ -135,7 +135,8 @@ def main(argv: Sequence[str]) -> None: ...@@ -135,7 +135,8 @@ def main(argv: Sequence[str]) -> None:
checkpoint = tf.train.Checkpoint(**checkpoint_dict) checkpoint = tf.train.Checkpoint(**checkpoint_dict)
checkpoint.restore(FLAGS.model_checkpoint).assert_existing_objects_matched() 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() model_for_serving.summary()
# TODO(b/194449109): Need to save the model to file and then convert tflite # TODO(b/194449109): Need to save the model to file and then convert tflite
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment