Allow export only path takes the hub module.

PiperOrigin-RevId: 322594580

Allow export only path takes the hub module.
PiperOrigin-RevId: 322594580
acaef198 · Hongkun Yu · A. Unique TensorFlower · 567bd18d · acaef198
Commit acaef198 authored Jul 22, 2020 by Hongkun Yu Committed by A. Unique TensorFlower Jul 22, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

official/nlp/bert/run_classifier.py official/nlp/bert/run_classifier.py +4 -1

No files found.
--- a/official/nlp/bert/run_classifier.py
+++ b/official/nlp/bert/run_classifier.py
@@ -343,7 +343,10 @@ def export_classifier(model_export_path, input_meta_data, bert_config,
  # Export uses float32 for now, even if training uses mixed precision.
  tf.keras.mixed_precision.experimental.set_policy('float32')
  classifier_model = bert_models.classifier_model(
-      bert_config, input_meta_data.get('num_labels', 1))[0]
+      bert_config,
+      input_meta_data.get('num_labels', 1),
+      hub_module_url=FLAGS.hub_module_url,
+      hub_module_trainable=False)[0]

  model_saving_utils.export_bert_model(
      model_export_path, model=classifier_model, checkpoint_dir=model_dir)