Skip to content

GitLab

Commit 62bf6fc0 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower
Browse files

Merge pull request #7957 from houtoms:ctl_add_persist_BN 

PiperOrigin-RevId: 287856865
parents 6fb5646f 0d1fb4cf
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
official/vision/image_classification/resnet_ctl_imagenet_main.py
View file @ 62bf6fc0
...@@ -199,6 +199,9 @@ def run(flags_obj): ...@@ -199,6 +199,9 @@ def run(flags_obj):
'mixed_bfloat16') 'mixed_bfloat16')
tf.compat.v2.keras.mixed_precision.experimental.set_policy(policy) 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. # TODO(anj-s): Set data_format without using Keras.
data_format = flags_obj.data_format data_format = flags_obj.data_format
if data_format is None: if data_format is None:
......
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