Skip to content

GitLab

Unverified Commit c07abee7 authored by Toby Boyd's avatar Toby Boyd Committed by GitHub
Browse files

remove turning off grappler layout optimizer. (#7384)

parent 8304e64a
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 5 deletions
official/transformer/v2/transformer_main.py
View file @ c07abee7
...@@ -127,7 +127,7 @@ class TransformerTask(object): ...@@ -127,7 +127,7 @@ class TransformerTask(object):
# We should have a better way in the tf.keras.mixed_precision API of doing # We should have a better way in the tf.keras.mixed_precision API of doing
# this. # this.
policy = tf.keras.mixed_precision.experimental.Policy( policy = tf.keras.mixed_precision.experimental.Policy(
'infer_float32_vars') "infer_float32_vars")
tf.keras.mixed_precision.experimental.set_policy(policy) tf.keras.mixed_precision.experimental.set_policy(policy)
def train(self): def train(self):
...@@ -135,9 +135,7 @@ class TransformerTask(object): ...@@ -135,9 +135,7 @@ class TransformerTask(object):
params, flags_obj, is_train = self.params, self.flags_obj, True params, flags_obj, is_train = self.params, self.flags_obj, True
# Sets config options. # Sets config options.
keras_utils.set_session_config( keras_utils.set_session_config(
enable_xla=flags_obj.enable_xla, enable_xla=flags_obj.enable_xla)
enable_grappler_layout_optimizer=
flags_obj.enable_grappler_layout_optimizer)
_ensure_dir(flags_obj.model_dir) _ensure_dir(flags_obj.model_dir)
if self.distribution_strategy: if self.distribution_strategy:
......
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