Merge pull request #7957 from houtoms:ctl_add_persist_BN

PiperOrigin-RevId: 287856865

Merge pull request #7957 from houtoms:ctl_add_persist_BN
PiperOrigin-RevId: 287856865
62bf6fc0 · A. Unique TensorFlower · 6fb5646f · 0d1fb4cf · 62bf6fc0
Commit 62bf6fc0 authored Jan 02, 2020 by A. Unique TensorFlower
Show whitespace changes
Inline Side-by-side

Showing with 3 additions and 0 deletions

official/vision/image_classification/resnet_ctl_imagenet_main.py ...l/vision/image_classification/resnet_ctl_imagenet_main.py +3 -0

No files found.
--- a/official/vision/image_classification/resnet_ctl_imagenet_main.py
+++ b/official/vision/image_classification/resnet_ctl_imagenet_main.py
@@ -199,6 +199,9 @@ def run(flags_obj):
        'mixed_bfloat16')
    tf.compat.v2.keras.mixed_precision.experimental.set_policy(policy)

+  # This only affects GPU.
+  common.set_cudnn_batchnorm_mode()
+
  # TODO(anj-s): Set data_format without using Keras.
  data_format = flags_obj.data_format
  if data_format is None: