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

PiperOrigin-RevId: 423714416

Author: renjie-liu

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