Internal change

PiperOrigin-RevId: 320559100

Internal change
PiperOrigin-RevId: 320559100
f5184e5f · André Susano Pinto · A. Unique TensorFlower · 0acca91f · f5184e5f
Commit f5184e5f authored Jul 10, 2020 by André Susano Pinto Committed by A. Unique TensorFlower Jul 10, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

official/nlp/bert/export_tfhub.py official/nlp/bert/export_tfhub.py +1 -1

No files found.
--- a/official/nlp/bert/export_tfhub.py
+++ b/official/nlp/bert/export_tfhub.py
@@ -79,7 +79,7 @@ def export_bert_tfhub(bert_config: configs.BertConfig,
                 do_lower_case, vocab_file)
  core_model, encoder = create_bert_model(bert_config)
  checkpoint = tf.train.Checkpoint(model=encoder)
-  checkpoint.restore(model_checkpoint_path).assert_consumed()
+  checkpoint.restore(model_checkpoint_path).assert_existing_objects_matched()
  core_model.vocab_file = tf.saved_model.Asset(vocab_file)
  core_model.do_lower_case = tf.Variable(do_lower_case, trainable=False)
  core_model.save(hub_destination, include_optimizer=False, save_format="tf")