Skip to content

GitLab

Commit dc93d9e5 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower
Browse files

Internal change 

PiperOrigin-RevId: 273562498
parent d0f21f25
Show whitespace changes
Inline Side-by-side
Showing with 6 additions and 3 deletions
official/vision/image_classification/resnet_imagenet_main.py
View file @ dc93d9e5
...@@ -232,6 +232,9 @@ def run(flags_obj): ...@@ -232,6 +232,9 @@ def run(flags_obj):
validation_freq=flags_obj.epochs_between_evals, validation_freq=flags_obj.epochs_between_evals,
verbose=2) verbose=2)
if flags_obj.enable_checkpoint_and_export: if flags_obj.enable_checkpoint_and_export:
if dtype == tf.bfloat16:
logging.warning("Keras model.save does not support bfloat16 dtype.")
else:
# Keras model.save assumes a float32 input designature. # Keras model.save assumes a float32 input designature.
export_path = os.path.join(flags_obj.model_dir, 'saved_model') export_path = os.path.join(flags_obj.model_dir, 'saved_model')
model.save(export_path, include_optimizer=False) model.save(export_path, include_optimizer=False)
......
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