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

PiperOrigin-RevId: 423714416
synwith · Jan 24, 2022 · 45848c4 · 45848c4
1 parent ab8b801
commit 45848c4
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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