"examples/roberta/vscode:/vscode.git/clone" did not exist on "1d44cc8520fc7d2fb4957fcab32d102de6b30626"
Commit 226f9419 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower
Browse files

Internal change

PiperOrigin-RevId: 462851231
parent 391f22eb
...@@ -71,9 +71,11 @@ class SavedModelCheckpointManager(tf.train.CheckpointManager): ...@@ -71,9 +71,11 @@ class SavedModelCheckpointManager(tf.train.CheckpointManager):
# Save the models for the checkpoint that just got written. # Save the models for the checkpoint that just got written.
saved_modules_directory = make_saved_modules_directory_name(checkpoint_path) saved_modules_directory = make_saved_modules_directory_name(checkpoint_path)
for model_name, model in self._modules_to_export.items(): for model_name, model in self._modules_to_export.items():
signatures = getattr(model, 'saved_model_signatures', None)
tf.saved_model.save( tf.saved_model.save(
obj=model, obj=model,
export_dir=os.path.join(saved_modules_directory, model_name)) export_dir=os.path.join(saved_modules_directory, model_name),
signatures=signatures)
saved_modules_directories_to_keep = [ saved_modules_directories_to_keep = [
make_saved_modules_directory_name(ckpt) for ckpt in self.checkpoints make_saved_modules_directory_name(ckpt) for ckpt in self.checkpoints
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment