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

Internal change

PiperOrigin-RevId: 437282849
parent e4be7e00
...@@ -74,8 +74,6 @@ def run(flags_obj): ...@@ -74,8 +74,6 @@ def run(flags_obj):
Returns: Returns:
Dictionary of training and eval stats. Dictionary of training and eval stats.
""" """
keras_utils.set_session_config(
enable_xla=flags_obj.enable_xla)
# Execute flag override logic for better model performance # Execute flag override logic for better model performance
if flags_obj.tf_gpu_thread_mode: if flags_obj.tf_gpu_thread_mode:
keras_utils.set_gpu_thread_mode_and_count( keras_utils.set_gpu_thread_mode_and_count(
...@@ -251,7 +249,8 @@ def run(flags_obj): ...@@ -251,7 +249,8 @@ def run(flags_obj):
optimizer=optimizer, optimizer=optimizer,
metrics=(['sparse_categorical_accuracy'] metrics=(['sparse_categorical_accuracy']
if flags_obj.report_accuracy_metrics else None), if flags_obj.report_accuracy_metrics else None),
run_eagerly=flags_obj.run_eagerly) run_eagerly=flags_obj.run_eagerly,
jit_compile=flags_obj.enable_xla)
train_epochs = flags_obj.train_epochs train_epochs = flags_obj.train_epochs
......
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