✍️ update tpu keras training script

Author: nglehuy

examples/conformer/train_tpu_keras_subword_conformer.py

@@ -92,11 +92,12 @@
 eval_dataset.load_max_lengths(args.max_lengths_prefix)
 
 with strategy.scope():
-    global_batch_size = config.learning_config.running_config.batch_size
+    batch_size = args.bs if args.bs is not None else config.learning_config.running_config.batch_size
+    global_batch_size = batch_size
     global_batch_size *= strategy.num_replicas_in_sync
     # build model
     conformer = Conformer(**config.model_config, vocabulary_size=text_featurizer.num_classes)
-    conformer._build(speech_featurizer.shape)
+    conformer._build(speech_featurizer.shape, prediction_shape=text_featurizer.prepand_shape, batch_size=batch_size)
     conformer.summary(line_length=120)
 
     optimizer = tf.keras.optimizers.Adam(