Internal change

PiperOrigin-RevId: 327984543

Internal change
PiperOrigin-RevId: 327984543
00c7ea79 · Hongkun Yu · A. Unique TensorFlower · 19e8d0a0 · 00c7ea79
Commit 00c7ea79 authored Aug 22, 2020 by Hongkun Yu Committed by A. Unique TensorFlower Aug 22, 2020
Show whitespace changes
Inline Side-by-side

Showing with 5 additions and 5 deletions

official/nlp/transformer/transformer_main.py official/nlp/transformer/transformer_main.py +5 -5

No files found.
--- a/official/nlp/transformer/transformer_main.py
+++ b/official/nlp/transformer/transformer_main.py
@@ -213,10 +213,10 @@ class TransformerTask(object):
        train_loss_metric = tf.keras.metrics.Mean(
            "training_loss", dtype=tf.float32)
        if params["enable_tensorboard"]:
-          summary_writer = tf.compat.v2.summary.create_file_writer(
-              flags_obj.model_dir)
+          summary_writer = tf.summary.create_file_writer(
+              os.path.join(flags_obj.model_dir, "summary"))
        else:
-          summary_writer = tf.compat.v2.summary.create_noop_writer()
+          summary_writer = tf.summary.create_noop_writer()
        train_metrics = [train_loss_metric]
        if params["enable_metrics_in_training"]:
          train_metrics = train_metrics + model.metrics
@@ -322,7 +322,7 @@ class TransformerTask(object):

          if params["enable_tensorboard"]:
            for metric_obj in train_metrics:
-              tf.compat.v2.summary.scalar(metric_obj.name, metric_obj.result(),
+              tf.summary.scalar(metric_obj.name, metric_obj.result(),
                                current_step)
              summary_writer.flush()