Remove vars dedup as keras fixed it.

PiperOrigin-RevId: 267525663

Remove vars dedup as keras fixed it.
PiperOrigin-RevId: 267525663
6d1dd03d · Hongkun Yu · A. Unique TensorFlower · 32974988 · 6d1dd03d
Commit 6d1dd03d authored Sep 05, 2019 by Hongkun Yu Committed by A. Unique TensorFlower Sep 05, 2019
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 3 deletions

official/bert/model_training_utils.py official/bert/model_training_utils.py +2 -3

No files found.
--- a/official/bert/model_training_utils.py
+++ b/official/bert/model_training_utils.py
@@ -231,9 +231,8 @@ def run_customized_training_loop(
      else:
        train_summary_writer = None

-      # De-dupes variables due to keras tracking issues.
-      training_vars = list({id(v): v for v in model.trainable_variables
-                           }.values())
+      # Collects training variables.
+      training_vars = model.trainable_variables

      def _replicated_step(inputs):
        """Replicated training step."""